Update setup_env.sh

pytorch · Nov 12, 2024 · a3dac35 · a3dac35
1 parent 1bfc967
commit a3dac35
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/userbenchmark/release-test/setup_env.sh b/userbenchmark/release-test/setup_env.sh
@@ -39,10 +39,8 @@ conda install -y -c pytorch ${MAGMA_VERSION}
 # install pip version of pytorch and torchvision
 if [[ ${PYTORCH_CHANNEL} == "pytorch-test" ]]; then
     pip3 install torch==${PYTORCH_VERSION} torchvision --index-url https://download.pytorch.org/whl/test/cu${CUDA_VERSION//./}
-    pip3 install jax
 else
     pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//./}
-    pip3 install jax
 fi
 
 python -c 'import torch; print(torch.__version__); print(torch.version.git_version)'