From e43be5074d920adc6bbc9a6fd3ff86bb05dd5178 Mon Sep 17 00:00:00 2001
From: Nicholas Geneva <5533524+NickGeneva@users.noreply.github.com>
Date: Wed, 22 Nov 2023 17:17:37 -0800
Subject: [PATCH] Fix graphcast jax dep in TOML (#124)

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 81c4aabb..11c038b2 100644
--- 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",
 ]