diff --git a/setup.py b/setup.py
index 8f7de74..d36c551 100644
--- a/setup.py
+++ b/setup.py
@@ -42,13 +42,14 @@
     package_data={"sbx": ["py.typed", "version.txt"]},
     install_requires=[
         "stable_baselines3>=2.8.0a0,<3.0",
-        "jax>=0.4.24,<0.7.0",  # tf probability not compatible yet with latest jax version
+        "jax>=0.4.24,<0.9.0",
         "jaxlib",
         "flax",
         "optax",
         "tqdm",
         "rich",
-        "tensorflow_probability",
+        # tf probability not compatible yet with latest jax version
+        "tfp-nightly>=0.26.0.dev20250831",
     ],
     extras_require={
         "tests": [