Update pyproject.toml

Fix jax dependency issue
2024-07-12 13:42:46 -07:00
parent 5b03a34927
commit 0a40da3c52
1 changed files with 1 additions and 1 deletions
@@ -5,7 +5,7 @@ name = "timesfm-jax"
 description = "Open weights time-series foundation model from Google Research."
 version = "1.0.1"
 dependencies = [
-    "jax==0.4.26",
+    "jax[cuda12]==0.4.26",
    "paxml==1.4.0",
    "praxis==1.4.0",
    "jaxlib==0.4.26",