From 5b2d629d11c31e93630cc1011bee2e6fb5c86f00 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 02:04:50 -0400
Subject: [PATCH] feat: test quantized tensor shapes directly

---
 test_quantize_shapes.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 test_quantize_shapes.py

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