Merge pull request #201 from apax-hub/unpin_jax

unpinned cuda jax version for Flax compatibility
apax-hub · Nov 14, 2023 · aadd2b4 · aadd2b4
2 parents 7f61d0d + ff0bed4
commit aadd2b4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -40,12 +40,12 @@ pip install --upgrade pip
 
 CUDA 12 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda12_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 CUDA 11 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda11_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 See the [Jax installation instructions](https://github.com/google/jax#installation) for more details.