change jax version

mlcommons · Jan 16, 2025 · 01eb881 · 01eb881
1 parent 1e62f15
commit 01eb881
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -128,10 +128,10 @@ jax_cpu =
 # JAX GPU
 # Note this installs both jax and jaxlib.
 jax_gpu =
-  jax==0.4.38
-  jaxlib==0.4.38
-  jax-cuda12-plugin[with_cuda]==0.4.38
-  jax-cuda12-pjrt==0.4.38
+  jax==0.4.36
+  jaxlib==0.4.36
+  jax-cuda12-plugin[with_cuda]==0.4.36
+  jax-cuda12-pjrt==0.4.36
   %(jax_core_deps)s
 
 # PyTorch CPU