Avoid KeyError

Lightning-AI · Nov 8, 2024 · a6b56ec · a6b56ec
1 parent a0a6151
commit a6b56ec
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py
@@ -137,7 +137,7 @@ def dce(trace: Trace, needed_proxies: None | set[Variable] = None) -> Trace:
         #   may mark some of the operation's outputs as unused
         some_unused = False
         for out in bsym.flat_proxy_outs:
-            if variableify(out) in needed_proxies and producer_map[out] == bsym:
+            if variableify(out) in needed_proxies and producer_map.get(out, None) == bsym:
                 needed = True
             else:
                 some_unused = True