cudnn11.7.1

syrkis · Jun 9, 2024 · 1267a80 · 1267a80
1 parent 8fcd56e
commit 1267a80
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -29,7 +29,8 @@ RUN python3.11 -m pip install -r requirements.txt
 
 # Install JAX with CUDA support. HPC is on CUDA 11, and JAX 0.2.25 is the latest version for that
 RUN python3.11 -m pip install --upgrade \
-    "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+    "jax[cuda11_pip]==0.4.25" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+    cuda11/jaxlib-0.4.25+cuda11.cudnn86-cp311-cp311-manylinux2014_x86_64.whl \
     optax
 
 # Set the environment variables