Fix CUDA graph replay: set device to cuda:0 before lm_head graph replay

2026-06-06 07:18:49 +00:00
parent f259d63930
commit 16b9a4def2
1 changed files with 2 additions and 1 deletions
--- 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