fix

Lightning-AI · Nov 13, 2024 · 4346842 · 4346842
1 parent b71214c
commit 4346842
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py
@@ -328,6 +328,8 @@ def __call__(self, *args, **kwargs):
             # vjp transform (applied later).
             def tag_tensorproxy_output_as_detached(proxy):
                 if isinstance(proxy, TensorProxy):
+                    # We need to remove name from trace, otherwise replace will return a proxy with new name.
+                    trace.names.remove(proxy.name)
                     return proxy.replace(tags=(ProxyTag.DETACHED_AUTOGRAD_GRAPH,))
                 return proxy