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

* Update trainer.py * format

Fix trainer saving safetensors: metadata is None (#28219)
* Update trainer.py * format
502a10a6 · hoshi-hiyouga · GitHub · cad9f5c6 · 502a10a6
Unverified Commit 502a10a6 authored Jan 02, 2024 by hoshi-hiyouga Committed by GitHub Jan 02, 2024
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

src/transformers/trainer.py src/transformers/trainer.py +3 -1

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -2908,7 +2908,9 @@ class Trainer:
            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: