uv installation instructions

2025-10-07 23:58:53 +00:00
parent 56e7fd67c6
commit 39df49c3a0
2 changed files with 9 additions and 68 deletions
@@ -23,12 +23,12 @@ torch = [
    "torch>=2.0.0",
 ]
 flax = [
-    "flax>=0.8.2",
-    "optax>=0.2.2",
-    "einshape>=0.8.0",
-    "orbax-checkpoint>=0.5.15",
-    "jaxtyping>=0.2.29",
-    "jax>=0.4.26",
+    "flax",
+    "optax",
+    "einshape",
+    "orbax-checkpoint",
+    "jaxtyping",
+    "jax[cuda12]"
 ]

 [tool.ruff]