From a6b56ec42d8c5ef96612b86e1a6187cb5c5c9188 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Fri, 8 Nov 2024 18:28:15 +0900 Subject: [PATCH] Avoid `KeyError` --- thunder/core/transform_common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py index e21ca4b28a..7b7a97469a 100644 --- 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