enhance rewrite state_dict missing _metadata (#14348)

This commit is contained in:
Chang Wang
2021-11-10 20:25:41 +08:00
committed by GitHub
parent 2b0d9389f8
commit bec02ff209

View File

@@ -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)