diff --git a/test_quantize_single_shard.py b/test_quantize_single_shard.py
index c235595..d2487f3 100644
--- a/test_quantize_single_shard.py
+++ b/test_quantize_single_shard.py
@@ -26,7 +26,7 @@ def test_single_shard(model_path, output_dir):
     shard_file = shards[0]
     print(f"Testing with single shard: {shard_file}")
     
-    state_dict = load_file(shard_file, device="cuda:0")
+    state_dict = load_file(shard_file, device="cpu")
     
     weight_keys = [
         k for k, v in state_dict.items()
@@ -40,9 +40,9 @@ def test_single_shard(model_path, output_dir):
     
     for key in weight_keys[:5]:  # Test first 5 layers
         try:
-            weight = state_dict[key]
+            weight = state_dict[key].to("cuda:0")
             qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
-            state_dict[key] = qweight
+            state_dict[key] = qweight.cpu()
             quantized += 1
             print(f"✓ {key}")
         except Exception as e: