diff --git a/setup.py b/setup.py
index 40fa8f82a7..5ff83fc31e 100755
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,7 @@
 tensorflow = [
     'dm-reverb==0.7.2',
     'keras==2.8.0',
-    'tensorflow-datasets==4.5.2',
+    'tensorflow-datasets',
     'tensorflow-estimator==2.8.0',
     'tensorflow==2.8.0',
     'tensorflow_probability==0.15.0',