"...git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "5d8eb93eeec0476b9f0fddc96f2960be0ce782b6"
Unverified Commit 502a10a6 authored by hoshi-hiyouga's avatar hoshi-hiyouga Committed by GitHub
Browse files

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

* Update trainer.py

* format
parent cad9f5c6
...@@ -2908,7 +2908,9 @@ class Trainer: ...@@ -2908,7 +2908,9 @@ class Trainer:
else: else:
logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
if self.args.save_safetensors: 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: else:
torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment