From 84df58091bdc5504dc7c22d79462d8d1b9cfb0e0 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Sat, 23 Dec 2023 15:39:21 +0800 Subject: [PATCH] Update trainer.py --- src/transformers/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 067d4ff3056018..15032a9b4c0f52 100755 --- 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: