diff --git a/setup.py b/setup.py
index 706e4c49..0f18fb99 100755
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,7 @@
 baselines_require = [
     'dm-sonnet',
     'dm-tree',
-    'tensorflow',
+    'tensorflow == 2.1',
     'tensorflow_probability >= 0.8, < 0.9',
     'trfl @ git+git://github.com/deepmind/trfl.git#egg=trfl',
     'tqdm',