Pinning jax version.

2024-05-21 20:02:57 +00:00
parent 469b7071dc
commit a4beacd023
5 changed files with 6 additions and 8 deletions
@@ -20,5 +20,5 @@ dependencies:
    - git+https://github.com/amazon-science/chronos-forecasting.git
    - praxis
    - paxml
-    - jax[cpu]
+    - jax[cpu]==0.4.26
    - einshape