diff --git a/pyproject.toml b/pyproject.toml
index 96696dd..d3e1e06 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ authors = [
     { name = "Martijn van den Ende", email = "martijnende@gmail.com" },
 ]
 dependencies = [
-    "jax==0.7.0",
+    "jax<=0.7.0",
     "jax_tqdm",
     "jaxtyping",
     "equinox",
@@ -22,7 +22,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["black", "isort", "twine"]
-gpu = ["jax[cuda12]==0.7.0"]
+gpu = ["jax[cuda12]<=0.7.0"]
 docs = [
     "make",
     "myst-nb",