From 0a40da3c522df2c9811a2293b63eddfdbf3d0fde Mon Sep 17 00:00:00 2001
From: Yichen Zhou <yichenzhou@google.com>
Date: Fri, 12 Jul 2024 13:42:46 -0700
Subject: [PATCH] Update pyproject.toml

Fix jax dependency issue
---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 91e1605..c4ec2e8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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",