feat: test quantized tensor shapes directly

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()