Change setup.cfg so that we install torch+cuda based on platform_mach…

…ine (architecture) not OS.
reginabarzilaygroup · May 31, 2024 · bc5c599 · bc5c599
1 parent 30b9ff3
commit bc5c599
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -46,11 +46,11 @@ install_requires =
     pydicom==2.3.0
     pylibjpeg[all]==2.0.0
     scikit-learn==1.0.2
-    torch==1.13.1+cu117; sys_platform != "darwin"
-    torch==1.13.1; sys_platform == "darwin"
+    torch==1.13.1+cu117; platform_machine == "x86_64"
+    torch==1.13.1; platform_machine != "x86_64"
     torchio==0.18.74
-    torchvision==0.14.1+cu117; sys_platform != "darwin"
-    torchvision==0.14.1; sys_platform == "darwin"
+    torchvision==0.14.1+cu117; platform_machine == "x86_64"
+    torchvision==0.14.1; platform_machine != "x86_64"
     tqdm==4.62.3
 
 [options.packages.find]