diff --git a/gpu.Dockerfile b/gpu.Dockerfile
index 62fc92fc..e467089a 100644
--- 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 && \