gpu cpu install (#81)

ergodicio · Oct 17, 2024 · 60e014e · 60e014e
1 parent f11e1c3
commit 60e014e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -36,7 +36,6 @@ def get_git_commit_hash():
     packages=find_packages(),
     python_requires=">=3.10",
     install_requires=[
-        "jax[cuda12]",
         "diffrax",
         "matplotlib",
         "scipy",
@@ -56,4 +55,8 @@ def get_git_commit_hash():
         "tabulate",
         "pydantic",
     ],
+    extras_require={
+        "cpu": ["jax"],
+        "gpu": ["jax[cuda12]"],
+    },
 )