cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: save as model.safetensors for from_pretrained

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 02:00:29 -04:00
parent 7115d56723
commit 1c792d23e4
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test_quantize_one_shard.py
View File

@@ -55,9 +55,9 @@ def test_one_shard(model_path, output_dir):
print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states") print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states")
# Save test output # Save test output as model.safetensors
state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()} state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
test_output = output_dir / "test_shard.safetensors" test_output = output_dir / "model.safetensors"
torch.save({**state_dict_cpu, **quant_states}, test_output) torch.save({**state_dict_cpu, **quant_states}, test_output)
print(f"Saved to: {test_output}") print(f"Saved to: {test_output}")