Update install_cuda.py

zackees · Jan 8, 2024 · 838108e · 838108e
1 parent b12631f
commit 838108e
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/install_cuda.py b/install_cuda.py
@@ -40,7 +40,16 @@
 
 # Delete the torch package if it doesn't have the cuda version
 if args.force or True:  # always do this path now.
-    subprocess.run(["pip", "uninstall", "-y", "torch", "torchtext", "torchdata", "torchaudio", "torchvision", "torch-directml"], check=True)
+    uninstall = [
+        "torch",
+        "torchtext",
+        "torchdata",
+        "torchaudio",
+        "torchvision",
+        "torch-directm"
+    ]
+    for package in uninstall:
+        subprocess.run(["pip", "uninstall", "-y", package], check=True)
     subprocess.run(["pip", "cache", "purge"], check=True)
 else:
     if TENSOR_CUDA_VERSION not in pip_list_stdout: