diff --git a/setup.py b/setup.py
index 20b8e10..f61c6bd 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
         "jax==0.4.14",
         "jaxlib==0.4.14",
         "objax==1.7.0",
-        "tensorflow_probability==0.20.1",
+        "tensorflow_probability==0.21",
         "numpy>=1.22"
     ],
     url='https://github.com/AaltoML/BayesNewton',