From 8fcd56e8e140b5f45deef0f8c3a9bdd830b26a2e Mon Sep 17 00:00:00 2001 From: Noah Syrkis Date: Sun, 9 Jun 2024 20:12:59 -0300 Subject: [PATCH] cudnn11.7.1 --- Dockerfile | 2 +- jobs/main.job | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index 86b3774..15d05cc 100644 --- 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 index eb428fa..aab4671 100644 --- 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