update install cuda script

zackees · Jan 7, 2024 · 9c12e5d · 9c12e5d
1 parent 9d7c6b1
commit 9c12e5d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/install_cuda.py b/install_cuda.py
@@ -14,10 +14,10 @@
 API_PY = os.path.join(HERE, "api.py")
 
 # Set the versions
-TENSOR_VERSION = "1.12.1"
-CUDA_VERSION = "cu116"
+TENSOR_VERSION = "2.1.2"
+CUDA_VERSION = "cu121"
 TENSOR_CUDA_VERSION = f"{TENSOR_VERSION}+{CUDA_VERSION}"
-EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cu116"
+EXTRA_INDEX_URL = f"https://download.pytorch.org/whl/{CUDA_VERSION}"
 
 
 argparser = argparse.ArgumentParser()