Updated _load_pretrained_model_low_mem to check if keys are in the state_dict (#16643)

* Updated _load_pretrained_model_low_mem to check if keys are in the stored state_dict * update after conversions

Updated _load_pretrained_model_low_mem to check if keys are in the state_dict (#16643)
* Updated _load_pretrained_model_low_mem to check if keys are in the stored state_dict * update after conversions
4099817b · Francesco Saverio Zuppichini · GitHub · 389f6615 · 4099817b
Unverified Commit 4099817b authored Apr 07, 2022 by Francesco Saverio Zuppichini Committed by GitHub Apr 07, 2022
Show whitespace changes
Inline Side-by-side

Showing with 8 additions and 7 deletions

src/transformers/modeling_utils.py src/transformers/modeling_utils.py +8 -7

No files found.
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -2162,6 +2162,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix

            # materialize state_dict entries one by one on CPU
            for k in loaded_state_dict_keys:
+                if k in state_dict:
                    submodule, param_name = find_submodule_and_param_name(model, k)
                    if submodule is not None:
                        param_dtype = getattr(submodule, param_name).dtype