diff --git a/quantize_streaming.py b/quantize_streaming.py
index 971a779..4ba362c 100644
--- a/quantize_streaming.py
+++ b/quantize_streaming.py
@@ -31,9 +31,9 @@ def streaming_quantize(model_path, output_path):
         print(f"Processing shard {shard_idx + 1}/{len(safetensors_files)}: {shard_file}")
         print(f"{'='*60}")
         
-        # Load shard to CPU
-        print("   Loading shard to CPU...")
-        shard_state_dict = torch.load(shard_file, map_location="cpu", weights_only=False)
+        # Load shard directly to GPU 0
+        print("   Loading shard to GPU 0...")
+        shard_state_dict = torch.load(shard_file, map_location="cuda:0", weights_only=False)
         
         # Quantize Linear layers in this shard using both GPUs
         print("   Quantizing Linear layers (both GPUs)...")
@@ -46,7 +46,7 @@ def streaming_quantize(model_path, output_path):
         gpu0_keys = weight_keys[::2]  # Even indices
         gpu1_keys = weight_keys[1::2]  # Odd indices
         
-        # Quantize on GPU 0
+        # Quantize on GPU 0 (shard already on GPU 0)
         print("      GPU 0: Quantizing...", end=" ")
         for key in gpu0_keys:
             tensor = shard_state_dict[key]
@@ -62,13 +62,17 @@ def streaming_quantize(model_path, output_path):
             )
             
             with torch.no_grad():
-                dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:0"))
+                dummy_linear.weight = nn.Parameter(tensor.clone())
                 _ = dummy_linear.weight.quant_state
             
-            shard_state_dict[key] = dummy_linear.weight.to("cpu")
+            shard_state_dict[key] = dummy_linear.weight
             quantized_keys += 1
         print(f"✓ {len(gpu0_keys)} layers")
         
+        # Move shard to GPU 1 for second half
+        print("      Moving to GPU 1...")
+        shard_state_dict = {k: v.to("cuda:1") if isinstance(v, torch.Tensor) else v for k, v in shard_state_dict.items()}
+        
         # Quantize on GPU 1
         print("      GPU 1: Quantizing...", end=" ")
         for key in gpu1_keys:
@@ -85,10 +89,10 @@ def streaming_quantize(model_path, output_path):
             )
             
             with torch.no_grad():
-                dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:1"))
+                dummy_linear.weight = nn.Parameter(tensor.clone())
                 _ = dummy_linear.weight.quant_state
             
-            shard_state_dict[key] = dummy_linear.weight.to("cpu")
+            shard_state_dict[key] = dummy_linear.weight
             quantized_keys += 1
         print(f"✓ {len(gpu1_keys)} layers")