From 16b9a4def224363fb1c37ef3429454106fae3f73 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 6 Jun 2026 07:18:49 +0000
Subject: [PATCH] Fix CUDA graph replay: set device to cuda:0 before lm_head
 graph replay

---
 single_shot_inference.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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