Don't require jax[cuda12] as the dep, just jax (#2)

Signed-off-by: Fabrice Normandin <[email protected]>

Author: lebrice

poetry.lock

@@ -8,7 +8,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.10"
-jax = {extras = ["cuda12"], version = "^0.4.28"}
+jax = "^0.4.28"
 torch = "^2.0.0"
 pytorch2jax = "^0.1.0"
 flax = "^0.8.4"
@@ -22,6 +22,7 @@ pre-commit = "^3.7.1"
 pytest-testmon = "^2.1.1"
 pytest-env = "^1.1.3"
 tensor-regression = "^0.0.4"
+jax = {extras = ["cuda12"], version = "^0.4.28"}
 
 
 [tool.poetry-dynamic-versioning]