diff --git a/test_quantize_shapes.py b/test_quantize_shapes.py new file mode 100644 index 0000000..457eaa8 --- /dev/null +++ b/test_quantize_shapes.py @@ -0,0 +1,29 @@ +#!/usr/bin/env python3 +"""Test quantized shard shapes.""" + +import torch +from pathlib import Path + + +def test_shapes(): + print("Loading quantized test shard...") + ckpt = torch.load("/data/models/test_quantize/model.safetensors", map_location="cpu") + + print(f"Loaded {len(ckpt)} tensors\n") + + # Check for mismatched shapes + mismatch_count = 0 + for key, tensor in list(ckpt.items())[:20]: # Check first 20 + if isinstance(tensor, torch.Tensor): + print(f"✓ {key} -> {tuple(tensor.shape)}") + elif hasattr(tensor, 'shape'): + print(f"✓ {key} -> {tensor.shape}") + else: + print(f"? {key} -> {type(tensor)}") + + print(f"\nTotal tensors: {len(ckpt)}") + print("If all shapes look correct, the quantization is working!") + + +if __name__ == "__main__": + test_shapes()