Fix graphcast jax dep in TOML (#124)

NVIDIA · Nov 23, 2023 · e43be50 · e43be50
1 parent 0f4b268
commit e43be50
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -59,7 +59,7 @@ pangu = [
     "onnxruntime-gpu>=1.15.1",
 ]
 graphcast = [
-    "jax[cuda12_pip]>=0.4.16",
+    "jax>=0.4.16",
     "graphcast @ https://github.com/deepmind/graphcast/archive/e622a15c1b9742d78f9b662f2af262604d58c204.tar.gz",
     "gcsfs>=2023.6.0",
 ]