[update] safe half

pull/825/head
jingyaogong 2026-01-30 20:29:31 +08:00
parent 3d2f03cfff
commit 9c1f4325f8
1 changed files with 2 additions and 3 deletions

View File

@ -64,9 +64,8 @@ def convert_torch2transformers_llama(torch_path, transformers_path, dtype=torch.
def convert_transformers2torch(transformers_path, torch_path): def convert_transformers2torch(transformers_path, torch_path):
model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
torch.save(model.state_dict(), torch_path) torch.save({k: v.cpu().half() for k, v in model.state_dict().items()}, torch_path)
print(f"模型已保存为 PyTorch 格式: {torch_path}") print(f"模型已保存为 PyTorch 格式 (half精度): {torch_path}")
if __name__ == '__main__': if __name__ == '__main__':