diff --git a/single_shot_inference.py b/single_shot_inference.py
index dbe385f0..54734263 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -1908,7 +1908,8 @@ def main():
             # Transfer last layer output to cuda:0 for lm_head graph
             graph_decoder.x_lm_in.copy_(X)
             
-            # lm_head graph replay
+            # lm_head graph replay — MUST be on cuda:0 (where it was captured)
+            torch.cuda.set_device(0)
             graph_decoder.lm_graph.replay()
             logits = graph_decoder.logits_buf