From bf23fb4880447ba19b208d31dc5a9d1ffda57487 Mon Sep 17 00:00:00 2001
From: Rajat Sen <rsen91@gmail.com>
Date: Wed, 8 Oct 2025 15:50:21 +0000
Subject: [PATCH] readme update and not pinning cuda versin for flax

---
 README.md      | 7 +++++++
 pyproject.toml | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

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]