diff --git a/requirements_mtj.txt b/requirements_mtj.txt
index d6a70923..5def7019 100644
--- a/requirements_mtj.txt
+++ b/requirements_mtj.txt
@@ -1,8 +1,8 @@
 numpy
 tqdm
 requests
-optax
-dm-haiku >= 0.0.5, <= 0.0.9
+optax >= 0.0.5, <= 0.0.9
+dm-haiku
 ray[default]
 jax == 0.2.12
 transformers