diff --git a/README.md b/README.md
index 9db8441..2d03cf9 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,13 @@ will be under construction over the next few weeks to
     uv pip install -e .[flax]
     ```
 
+3. [Optional] Install your preferred `torch` / `jax` backend based on your OS and accelerators
+(CPU, GPU, TPU or Apple Silicon).:
+
+-   [Install PyTorch](https://pytorch.org/get-started/locally/).
+-   [Install Jax](https://docs.jax.dev/en/latest/installation.html#installation)
+    for Flax.
+
 ### Code Example
 
 ```python
diff --git a/pyproject.toml b/pyproject.toml
index 36ce3fd..facfbdb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ flax = [
     "einshape",
     "orbax-checkpoint",
     "jaxtyping",
-    "jax[cuda12]"
+    "jax[cuda]"
 ]
 
 [tool.ruff]