use original loaded keys to find mismatched keys (#16920)

2d91e3c3 · Yongliang Shen · GitHub · d365f507 · 2d91e3c3
Unverified Commit 2d91e3c3 authored Apr 27, 2022 by Yongliang Shen Committed by GitHub Apr 26, 2022
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

src/transformers/modeling_utils.py src/transformers/modeling_utils.py +3 -2

No files found.
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -2022,6 +2022,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
                return key.replace("gamma", "weight")
            return key

+        original_loaded_keys = loaded_keys
        loaded_keys = [_fix_key(key) for key in loaded_keys]

        if len(prefix) > 0:
@@ -2114,7 +2115,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
            mismatched_keys = _find_mismatched_keys(
                state_dict,
                model_state_dict,
-                loaded_keys,
+                original_loaded_keys,
                add_prefix_to_model,
                remove_prefix_from_model,
                ignore_mismatched_sizes,
@@ -2140,7 +2141,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
                mismatched_keys += _find_mismatched_keys(
                    state_dict,
                    model_state_dict,
-                    loaded_keys,
+                    original_loaded_keys,
                    add_prefix_to_model,
                    remove_prefix_from_model,
                    ignore_mismatched_sizes,