diff --git a/single_shot_inference.py b/single_shot_inference.py
index 07571807..1fb8ed80 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -366,9 +366,10 @@ def main():
             target_gpu = li % NUM_GPUS
             target_device = f"cuda:{target_gpu}"
             
-            # Move activation to layer's GPU
+            # Move activation to layer's GPU and set as active device
             if x.device != torch.device(target_device):
                 x = x.to(target_device)
+            torch.cuda.set_device(target_gpu)
             
             lw = layer_weights[li]
             rc, rs = rope_caches[target_gpu]