diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 04f61d9d4..1148aac67 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1049,7 +1049,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: - state_dict = {k: v for k, v in state_dict.items() if k not in self._keys_to_ignore_on_save} + for ignore_key in self._keys_to_ignore_on_save: + 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)