diff --git a/tensor2tensor/utils/t2t_model.py b/tensor2tensor/utils/t2t_model.py
index c4ec2b87e..21e3aefe0 100644
--- a/tensor2tensor/utils/t2t_model.py
+++ b/tensor2tensor/utils/t2t_model.py
@@ -1337,7 +1337,7 @@ def estimator_model_fn(cls,
       # by logits["self_generated_targets"].
       tf.logging.info("Replacing targets with model-provided targets.")
       features["targets"] = labels = logits.pop("self_generated_targets")
-      assert logits.keys() == ["logits"], (
+      assert list(logits.keys()) == ["logits"], (
           # See "Returns" in the "top" method docstring for the expected
           # "logits" format when targets are generated at training time.
           "Expect only key 'logits' when there is 'self_generated_targets'. "