Filter out nonproxies in find_producer_symbols

Lightning-AI · May 3, 2024 · 14cd04e · 14cd04e
1 parent ba7f4fd
commit 14cd04e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/thunder/core/utils.py b/thunder/core/utils.py
@@ -1048,6 +1048,7 @@ def find_producer_symbols(trace: TraceCtx, proxies: Sequence[Proxy], stop_proxie
         (__b = ltorch.sub(x, y)
         # __b = prims.sub(x, y),)
     """
+    stop_proxies = filter(lambda x: isinstance(x, Proxy), stop_proxies)
     trace_producers = producers(trace)
     result = set()
     queue = list(proxies)