fix: free GPU memory after each quantization

test_quantize_single_shard.py

@@ -38,16 +38,21 @@ def test_single_shard(model_path, output_dir):
     quantized = 0
     failed = 0
     
-    for key in weight_keys[:5]:  # Test first 5 layers
+    for i, key in enumerate(weight_keys[:5]):
         try:
             weight = state_dict[key].to("cuda:0")
             qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
             state_dict[key] = qweight.cpu()
+            del weight, qweight
+            gc.collect()
+            torch.cuda.empty_cache()
             quantized += 1
             print(f"✓ {key}")
         except Exception as e:
             failed += 1
             print(f"✗ {key}: {e}")
+            gc.collect()
+            torch.cuda.empty_cache()
     
     print(f"\nResults: {quantized} quantized, {failed} failed")