Fix trainer saving safetensors: metadata is None (huggingface#28219)

* Update trainer.py * format
epfl-dlab · Jan 4, 2024 · 62911a3 · 62911a3
1 parent 8d85c88
commit 62911a3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2908,7 +2908,9 @@ 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: