Dynamo ignores frame using yield

[xmfan]

🐛 Describe the bug

OptimizeContext is wrapped, and frame evaluation callback seems to be set to convert_frame. But there's no dynamo logs and the entire frame seems to run in eager.

import torch
from torch._dynamo import config

def fn():
    model = torch.nn.Sequential(
        torch.nn.Linear(400, 4),
        torch.nn.Sigmoid(),
    )

    for _ in range(3):
        x = torch.randn([100, 400])

        result = model(x).sum()
        result.backward()
        yield model[0].weight.grad
        yield model[0].bias.grad
        model.zero_grad()

opt_fn = torch.compile(fn)
outs = opt_fn()
res = list(outs)

Versions

main

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[ezyang]

If you wrap model in a user callable I think it will work