fix training issues (#36158)

* fix training issues * Update Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
huggingface · Feb 13, 2025 · 0ca7259 · 0ca7259
1 parent 845b0a2
commit 0ca7259
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/transformers/models/paligemma/configuration_paligemma.py b/src/transformers/models/paligemma/configuration_paligemma.py
@@ -74,6 +74,7 @@ class PaliGemmaConfig(PretrainedConfig):
 
     model_type = "paligemma"
     sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
+    keys_to_ignore_at_inference = ["past_key_values"]
 
     def __init__(
         self,

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -4477,7 +4477,7 @@ def prediction_step(
         inputs = self._prepare_inputs(inputs)
         if ignore_keys is None:
             if hasattr(self.model, "config"):
-                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
+                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", ["past_key_values"])
             else:
                 ignore_keys = []