diff --git a/setup.py b/setup.py
index e203237d6c..81574096aa 100644
--- a/setup.py
+++ b/setup.py
@@ -49,6 +49,7 @@
     "tensorflow-cpu>=2.4.0",
     "tensorflow_text>=2.4.0",  # WMT example.
     "tensorflow_datasets",
+    "tensorflow==2.4.1",  # TODO(marcvanzee): Remove once #1326 is fixed.
 ]
 
 __version__ = None