add extra_require for jax[cuda12]

adaptive-intelligent-robotics · Aug 26, 2024 · 380cc81 · 380cc81
1 parent 86c7f9a
commit 380cc81
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -35,6 +35,9 @@
         "scikit-learn>=1.0.2",
         "scipy>=1.8.0",
     ],
+    extras_require={
+        "cuda12": ["jax[cuda12]>=0.4.16"],
+    },
     dependency_links=[
         "https://storage.googleapis.com/jax-releases/jax_releases.html",
     ],