Update torch_package.py

Related to #214
gitmylo · Feb 12, 2024 · 944cd19 · 944cd19
1 parent 240bbbd
commit 944cd19
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup_tools/magicinstaller/requirements/torch_package.py b/setup_tools/magicinstaller/requirements/torch_package.py
@@ -14,6 +14,6 @@ def is_installed(self):
 
     def install(self):
         # if self.is_windows():
-            return self.install_pip('torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118', 'PyTorch')
+            return self.install_pip('torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --extra-index-url https://download.pytorch.org/whl/cu118', 'PyTorch')
         # else:
         #     return self.install_pip('torch torchvision torchaudio', 'PyTorch')