diff --git a/src/transformers/trainer_utils.py b/src/transformers/trainer_utils.py
index afc2e0d1561c..46fd0cdd05b6 100644
--- a/src/transformers/trainer_utils.py
+++ b/src/transformers/trainer_utils.py
@@ -676,6 +676,8 @@ def __init__(
         self.message_logged = False
 
     def _remove_columns(self, feature: dict) -> dict:
+        if not isinstance(feature, dict):
+            return feature
         if not self.message_logged and self.logger and self.model_name:
             ignored_columns = list(set(feature.keys()) - set(self.signature_columns))
             if len(ignored_columns) > 0: