change back

divyashreepathihalli · divyashreepathihalli · commit 9d012c471203 · 2025-02-20T23:22:44.000Z
diff --git a/requirements.txt b/requirements.txt
@@ -10,7 +10,7 @@ torch==2.6.0+cpu
 torch-xla==2.6.0;sys_platform != 'darwin'
 
 # Jax.
-jax[cpu]>=0.4.31
+jax[cpu]
 flax
 
 # Common deps.
