fix jax version

oxcsml · Jul 20, 2022 · 19627b5 · 19627b5
1 parent 85d8c3e
commit 19627b5
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,8 +1,5 @@
 --find-links https://storage.googleapis.com/jax-releases/jax_releases.html
-# jax[cuda11_cudnn805]
-# jaxlib==0.1.74+cuda11.cudnn805  # not sure that version will work on other zizgpu0x
-jax==0.3.1+cuda11.cudnn805
-# jaxlib==0.3.0+cuda11.cudnn805  # not sure that version will work on other zizgpu0x
+jax[cuda11_cudnn805]==0.3.1
 dm-haiku
 optax
 fsspec[http]>=2021.05.0, !=2021.06.0