From cee50021712a8e32f81e6f501174193068d531b2 Mon Sep 17 00:00:00 2001
From: Charles Xu <92070992+charlesxu0124@users.noreply.github.com>
Date: Fri, 14 Jun 2024 13:56:31 -0700
Subject: [PATCH] Update jax cuda installation instructions

---
 README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 509aade..b05f3ea 100644
--- a/README.md
+++ b/README.md
@@ -37,9 +37,9 @@ We fixed a major issue in the intervention action frame. See release [v0.1.1](ht
         pip install --upgrade "jax[cpu]"
         ```
 
-    - For GPU: (change cuda12 to cuda11 if you are using older driver versions)
+    - For GPU:
         ```bash
-        pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+        pip install -U "jax[cuda12]"
         ```
 
     - For TPU