Pinning jax version.

2024-05-21 20:02:57 +00:00
parent 469b7071dc
commit a4beacd023
5 changed files with 6 additions and 8 deletions
@@ -1,4 +1,4 @@
-name: tfm_env
+name: tfm_env_t2

 channels:
  - conda-forge
@@ -13,6 +13,5 @@ dependencies:
    - utilsforecast
    - praxis
    - paxml
-    - -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  
-    - "jax[cuda12_pip]"
+    - jax[cuda12]==0.4.26
    - einshape