diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
index 38a4e8cc10512d..cd8e6956448778 100644
--- a/src/transformers/modeling_tf_utils.py
+++ b/src/transformers/modeling_tf_utils.py
@@ -1279,7 +1279,7 @@ def serving_output(self, output):
         """
         if not isinstance(output, ModelOutput):
             return output
-        for key in output.keys():
+        for key in output:
             if key.endswith("hidden_states") and not getattr(self.config, "output_hidden_states", False):
                 output[key] = None
             elif key.endswith("attentions") and not getattr(self.config, "output_attentions", False):