Update setup_env.sh

pytorch · Apr 8, 2024 · aa413c8 · aa413c8
1 parent 82b465e
commit aa413c8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/userbenchmark/release-test/setup_env.sh b/userbenchmark/release-test/setup_env.sh
@@ -44,13 +44,13 @@ conda install -y -c pytorch ${MAGMA_VERSION}
 
 # install pip version of pytorch and torchvision
 if [[ ${PYTORCH_CHANNEL} == "pytorch-test" ]]; then
-    pip3 install --pre torch==2.2.2 --extra-index-url https://download.pytorch.org/whl/test/cu121
-    pip3 install --pre torchvision==0.17.2 torchaudio==2.2.2 --extra-index-url https://download.pytorch.org/whl/test/cu121
+    pip3 install --pre torch==2.3.0 --extra-index-url https://download.pytorch.org/whl/test/cu121
+    pip3 install --pre torchvision==0.18.0 torchaudio==2.3.0 --extra-index-url https://download.pytorch.org/whl/test/cu121
     #conda install -y pytorch=2.1.2 torchvision=0.16.2 torchaudio=2.1.2 pytorch-cuda=12.1 -c pytorch-test -c nvidia
     #pip3 install --force-reinstall --no-cache-dir  torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/test/cu${CUDA_VERSION//./}
 else
     pip3 install --force-reinstall --no-cache-dir torch torchvision torchaudio
-    #conda install -y pytorch=2.1.1 torchvision=0.16.1 torchaudio=2.1.1 pytorch-cuda=12.1 -c pytorch -c nvidia
+    #conda install -y pytorch=2.2.2 torchvision=0.17.2 torchaudio=2.2.2 pytorch-cuda=12.1 -c pytorch -c nvidia
 fi
 
 python -c 'import torch; print(torch.__version__); print(torch.version.git_version)'