Update JAX CUDA release

iree-org · Apr 8, 2024 · 208593d · 208593d
1 parent 7a2f9c4
commit 208593d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/setup_venv.sh b/common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/setup_venv.sh
@@ -37,7 +37,7 @@ if [ -z "$WITH_CUDA" ]; then
   python -m pip install --upgrade "jax[cpu]" "flax"
 else
   echo "Installing jax and dependencies with cuda support"
-  python -m pip install --upgrade "jax[cuda11_local]" "flax" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+  python -m pip install --upgrade "jax[cuda11]" "flax" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 fi
 
 # Install Tensorflow to generate TFLite models.

diff --git a/comparative_benchmark/jax/setup_venv.sh b/comparative_benchmark/jax/setup_venv.sh
@@ -31,7 +31,7 @@ if [ -z "$WITH_CUDA" ]; then
   python -m pip install --upgrade "jax[cpu]" "flax"
 else
   echo "Installing jax and dependencies with cuda support"
-  python -m pip install --upgrade "jax[cuda11_local]" "flax" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+  python -m pip install --upgrade "jax[cuda11]" "flax" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 fi
 
 python -m pip install --upgrade -r "${TD}/requirements.txt"