Fix dependency error

pomonam · Apr 9, 2024 · 2dc9db5 · 2dc9db5
1 parent 617fcaa
commit 2dc9db5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -43,14 +43,14 @@ python_requires = >=3.8
 [options.extras_require]
 # JAX CPU
 jax_cpu =
-  jax==0.4.26
-  jaxlib==0.4.26
+  jax==0.4.13
+  jaxlib==0.4.13
   dm-haiku==0.0.12
 
 # JAX GPU
 jax_gpu =
-  jax[cuda]==0.4.26
-  jaxlib==0.4.264+cuda11.cudnn82
+  jax[cuda]==0.4.13
+  jaxlib==0.4.13+cuda11.cudnn82
   dm-haiku==0.0.12
 
 ###############################################################################