Add cublas 12.1.0.26

pytorch · Dec 1, 2023 · 1173d5f · 1173d5f
1 parent a98fe7c
commit 1173d5f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/userbenchmark/release-test/setup_env.sh b/userbenchmark/release-test/setup_env.sh
@@ -38,7 +38,8 @@ conda install -y -c pytorch ${MAGMA_VERSION}
 
 # install pip version of pytorch and torchvision
 if [[ ${PYTORCH_CHANNEL} == "pytorch-test" ]]; then
-    pip3 install torch torchvision --index-url https://download.pytorch.org/whl/test/cu121
+    pip3 install torch torchvision --index-url https://download.pytorch.org/whl/test/cu${CUDA_VERSION//./}
+    pip3 install nvidia-cublas-cu12==12.1.0.26
 else
     pip3 install torch torchvision
 fi

diff --git a/utils/cuda_utils.py b/utils/cuda_utils.py
@@ -45,12 +45,6 @@ def prepare_cuda_env(cuda_version: str, dryrun=False):
     env["LD_LIBRARY_PATH"] = f"{cuda_path_str}/lib64:{cuda_path_str}/extras/CUPTI/lib64:{env['LD_LIBRARY_PATH']}"
     if dryrun:
         print(f"CUDA_HOME is set to {env['CUDA_HOME']}")
-
-     # step 1.5: install the correct nvcc version
-    if cuda_version ==  "12.1":
-        install_nvcc = ["pip", "install", "nvidia-cuda-nvcc-cu12"]
-        subprocess.check_call(install_nvcc, env=env)
-
     # step 2: test call to nvcc to confirm the version is correct
     test_nvcc = ["nvcc", "--version"]
     if dryrun:
@@ -59,7 +53,6 @@ def prepare_cuda_env(cuda_version: str, dryrun=False):
         output = subprocess.check_output(test_nvcc, stderr=subprocess.STDOUT, env=env).decode()
         print(f"NVCC version output: {output}")
         assert _nvcc_output_match(output, cuda_version), f"Expected CUDA version {cuda_version}, getting nvcc test result {output}"
-
     # step 3: install the correct magma version
     install_magma_cmd = ["conda", "install", "-y", "-c", "pytorch", CUDA_VERSION_MAP[cuda_version]['magma_version']]
     if dryrun: