SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

enhance rewrite state_dict missing _metadata (#14348)

Browse Source
This commit is contained in:
Chang Wang
2021-11-10 20:25:41 +08:00
committed by GitHub
parent 2b0d9389f8
commit bec02ff209
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/modeling_utils.py
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 # Handle the case where some state_dict keys shouldn't be saved
if self._keys_to_ignore_on_save is not None: if self._keys_to_ignore_on_save is not None:
for ignore_key in self._keys_to_ignore_on_save: 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` # If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(save_directory, WEIGHTS_NAME) output_model_file = os.path.join(save_directory, WEIGHTS_NAME)