From bec02ff209d20d4bca974dab99a8c4e11f5fa737 Mon Sep 17 00:00:00 2001 From: Chang Wang Date: Wed, 10 Nov 2021 20:25:41 +0800 Subject: [PATCH] enhance rewrite state_dict missing _metadata (#14348) --- src/transformers/modeling_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 1148aac67d..213f5b57ee 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1050,7 +1050,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix # Handle the case where some state_dict keys shouldn't be saved if self._keys_to_ignore_on_save is not None: for ignore_key in self._keys_to_ignore_on_save: - del state_dict[ignore_key] + if ignore_key in state_dict.keys(): + del state_dict[ignore_key] # If we save using the predefined names, we can load using `from_pretrained` output_model_file = os.path.join(save_directory, WEIGHTS_NAME)