Update trainer.py

huggingface · Dec 23, 2023 · 84df580 · 84df580
1 parent 29e7a1e
commit 84df580
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2908,7 +2908,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
             else:
                 logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                 if self.args.save_safetensors:
-                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME))
+                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"})
                 else:
                     torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
         else: