fix: save as model.safetensors for from_pretrained

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}")