Merge pull request #66 from google-research/rajat_dev

minor changes to environment files
2024-06-03 15:33:37 -07:00
parent 67041232c4 1b267eb70c
commit 7d97953b34
2 changed files with 2 additions and 2 deletions
@@ -20,7 +20,7 @@ dependencies:
    - git+https://github.com/amazon-science/chronos-forecasting.git
    - praxis
    - paxml
-    - jax[cuda12]
+    - jax[cuda12]==0.4.26
    - einshape
    - python-dotenv
    - nixtla>=0.5.1
@@ -20,7 +20,7 @@ dependencies:
    - git+https://github.com/amazon-science/chronos-forecasting.git
    - praxis
    - paxml
-    - jax[cpu]
+    - jax[cpu]==0.4.26
    - einshape
    - python-dotenv
    - nixtla>=0.5.1