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