Update install_cuda.py

zackees · Jan 8, 2024 · 7799595 · 7799595
1 parent e5b829d
commit 7799595
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/install_cuda.py b/install_cuda.py
@@ -41,6 +41,8 @@
 # Delete the torch package if it doesn't have the cuda version
 if args.force:
     subprocess.run(["pip", "uninstall", "-y", "torch"], check=True)
+    subprocess.run(["pip", "uninstall", "-y", "torchvision"], check=True)
+    subprocess.run(["pip", "uninstall", "-y", "torch-directml"], check=True)
     subprocess.run(["pip", "cache", "purge"], check=True)
 else:
     if TENSOR_CUDA_VERSION not in pip_list_stdout: