Upgrade JAX to 0.2.16 (#1039)

Kaggle · Jul 2, 2021 · 882546e · 882546e
1 parent 7e3c7a2
commit 882546e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/gpu.Dockerfile b/gpu.Dockerfile
@@ -78,8 +78,8 @@ RUN pip uninstall -y lightgbm && \
     /tmp/clean-layer.sh
 
 # Install JAX (Keep JAX version in sync with CPU image)
-RUN  pip install jax==0.2.12 jaxlib==0.1.64+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
-     /tmp/clean-layer.sh
+RUN pip install jax==0.2.16 jaxlib==0.1.68+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
+    /tmp/clean-layer.sh
 
 # Reinstall packages with a separate version for GPU support.
 RUN pip uninstall -y tensorflow && \