cudnn11.7.1

syrkis · Jun 9, 2024 · 8fcd56e · 8fcd56e
1 parent 7094ca2
commit 8fcd56e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -29,7 +29,7 @@ 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]==0.4.25" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+    "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
     optax
 
 # Set the environment variables

diff --git a/jobs/main.job b/jobs/main.job
@@ -11,6 +11,7 @@
 
 
 module --ignore-cache load singularity/3.4.1
-module --ignore-cache load CUDA/11.1.1-GCC-10.2.0
+module --ignore-cache load CUDA/11.7.0
+module --ignore-cache load cuDNN/8.4.1.50-CUDA-11.7.0
 
 srun singularity exec --nv container.sif python3.11 main.py