From 1c792d23e4fb9875f6e2c69c9237e01b59ce3c43 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 02:00:29 -0400
Subject: [PATCH] fix: save as model.safetensors for from_pretrained

---
 test_quantize_one_shard.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test_quantize_one_shard.py b/test_quantize_one_shard.py
index 43921e5..b147a0b 100644
--- a/test_quantize_one_shard.py
+++ b/test_quantize_one_shard.py
@@ -55,9 +55,9 @@ def test_one_shard(model_path, output_dir):
     
     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()}
-    test_output = output_dir / "test_shard.safetensors"
+    test_output = output_dir / "model.safetensors"
     torch.save({**state_dict_cpu, **quant_states}, test_output)
     print(f"Saved to: {test_output}")