Merge pull request #2082 from atalman/more_testing212

test
pytorch · Dec 12, 2023 · 631e944 · 631e944
2 parents 1c7cd09 + ad01287
commit 631e944
Showing 1 changed file with 2 additions and 1 deletion.
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==2.1.2 torchvision --index-url https://download.pytorch.org/whl/test/cu${CUDA_VERSION//./}
+    pip3 install --force-reinstall  torch==2.1.2 torchvision=0.16.2 torchaudio=2.1.2 --index-url https://download.pytorch.org/whl/test/cu${CUDA_VERSION//./}
+    # pip3 install --force-reinstall torch==2.1.2 torchvision
 else
     pip3 install --force-reinstall torch torchvision
 fi