fix output tuple wrapping

Lightning-AI · Jun 24, 2024 · 083a3b1 · 083a3b1
1 parent aeebb1b
commit 083a3b1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/thunder/core/jit_ext.py b/thunder/core/jit_ext.py
@@ -1690,13 +1690,16 @@ def thunder_general_jit(
     pro_to_comp = tuple(sorted(pro_to_comp, key=lambda v: proxy_order[id(v.proxy)]))
 
     bind_inputs("computation", computation_trace, pro_to_comp_proxies)
+
     if epilogue_trace:
         l = len(epilogue_trace.bound_symbols)
         if l == 0:
             epilogue_trace = None
         elif l == 1:
             (r,) = epilogue_trace.bound_symbols
             assert r.sym == prims.python_return
+            assert computation_trace.bound_symbols[-1].sym == prims.python_return
+            computation_trace.bound_symbols[-1].args = r.args
             epilogue_trace = None
 
     if epilogue_trace: