updated jax version

KosinskiLab · Feb 20, 2024 · 565fe55 · 565fe55
1 parent 3ea89dd
commit 565fe55
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/workflow/docker/pulldown.dockerfile b/workflow/docker/pulldown.dockerfile
@@ -14,7 +14,7 @@
 
 # Extended by Valentin Maurer <[email protected]>
 
-ARG CUDA=11.8
+ARG CUDA=11.1.1
 FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu18.04
 ARG CUDA
 
@@ -69,9 +69,9 @@ RUN pip3 install --upgrade pip --no-cache-dir \
     && pip3 install --upgrade --no-cache-dir \
       alphapulldown==1.0.4 \
       pytest \
-      jax==0.4.23 \
-      jaxlib==0.4.23 \
-      jaxlib==0.4.23+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+      jax==0.3.25 \
+      jaxlib==0.3.25+cuda11.cudnn805 \
+      -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 RUN chmod u+s /sbin/ldconfig.real