diff --git a/env/hpc/requirements.txt b/env/hpc/requirements.txt
index 5cdc7f1..d9527d7 100644
--- a/env/hpc/requirements.txt
+++ b/env/hpc/requirements.txt
@@ -4,7 +4,7 @@ evosax==0.2.0
 flax>=0.12.2
 gymnasium>=1.2.3
 ipykernel==7.2.0
-jax==0.9.0.1
+jax[cuda12]==0.9.0.1
 numpy>=2.0.0
 scipy>=1.14.0
 protobuf>=5.0.0
diff --git a/pyproject.toml b/pyproject.toml
index 556593b..d1d468d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
 
 [project.optional-dependencies]
 cuda = [
-    "jax[cuda13]==0.9.0.1",
+    "jax[cuda12]==0.9.0.1",
 ]
 
 [dependency-groups]