[dynamo] KeyError in resume function creation in Python 3.11+ #162313
Description
Minimal repro:
import torch
flag = True
@torch.compile(backend="eager")
def fn(x):
x = x + 1
torch._dynamo.graph_break()
x = x + 2
if flag:
with torch.no_grad():
torch._dynamo.graph_break()
else:
with torch.no_grad():
torch._dynamo.graph_break()
return x + 4
fn(torch.ones(3))
flag = False
fn(torch.ones(3))This is because we're incorrectly using meta.block_target_offset_remap in resume_execution.py.
Internal xref: T236677438
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @Lucaskabela