maskrcnn.py

import torch
import torchvision

import torch._dynamo
torch._dynamo.config.assume_static_by_default = True

device = "cuda:0"

model = torchvision.models.detection.maskrcnn_resnet50_fpn()
model = model.to(device)
model = torch.compile(model, dynamic=True)
model.eval()

# This works
for _ in range(1):
    x = torch.rand(1, 3, 800, 800, device=device)
    r = model(x)
    r[0]['boxes'].sum().backward()
    #next(iter(r[0].items()))[0].sum().backward()

# This does not
#with torch.no_grad():
#    _ = model(x)