diff --git a/bsuite/_metadata.py b/bsuite/_metadata.py
index 598cbcc..f875b3b 100644
--- a/bsuite/_metadata.py
+++ b/bsuite/_metadata.py
@@ -20,4 +20,4 @@
 a time when bsuite's dependencies may not have been installed yet.
 """
 
-__version__ = '0.3.4'
+__version__ = '0.3.5'
diff --git a/setup.py b/setup.py
index d4b203c..e317829 100755
--- a/setup.py
+++ b/setup.py
@@ -25,8 +25,8 @@
 baselines_require = [
     'dm-sonnet',
     'dm-tree',
-    'tensorflow == 2.1',
-    'tensorflow_probability >= 0.8, < 0.9',
+    'tensorflow',
+    'tensorflow_probability',
     'trfl',
     'tqdm',
 ]