gistfile1.txt

import torch

def compiler_fn(gm):
    return torch.compile(gm, mode="reduce-overhead", fullgraph=True, dynamic=True)

def fn():
    x = torch.randn(2, 2, device="cuda", requires_grad=True)
    y = torch.randn(2, 2, device="cuda")
    out = torch.mm(x, y)
    loss = out.sum() / out.numel()
    loss.backward()

with torch._dynamo.compiled_autograd.enable(compiler_fn):
    fn()