improve rewrite state_dict missing _metadata (#14276)

huggingface · Nov 16, 2021 · c8206b4 · c8206b4
1 parent b6b97c3
commit c8206b4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -1049,7 +1049,8 @@ def save_pretrained(
 
         # 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)