Restrict jaxlib version in TPU Colabs

2022-07-12 16:30:26 -04:00 · 2022-07-12 16:30:26 -04:00 · 70aa182671
parent dd6da50e58
commit 70aa182671
1 changed files with 1 additions and 0 deletions
--- a/requirements_mtj.txt
+++ b/requirements_mtj.txt
@ -5,6 +5,7 @@ requests
 optax >= 0.0.5, <= 0.0.9
 dm-haiku == 0.0.5
 jax == 0.2.21
+jaxlib >= 0.1.69, <= 0.3.7
 transformers >= 4.19
 progressbar2
 git+https://github.com/VE-FORBRYDERNE/mesh-transformer-jax@ck