importtorchmodel=torch.nn.Sequential(*[torch.nn.Linear(1024,1024,False,device="cuda")for_inrange(10)])input=torch.rand(1024,device="cuda")output=model(input)output.sum().backward()

# exit cleanly if we are on a device that doesn't support torch.compileiftorch.cuda.get_device_capability()<(7,0):print("Exiting because torch.compile is not supported on this device.")importsyssys.exit(0)opt=torch.optim.Adam(model.parameters(),lr=0.01)@torch.compile(fullgraph=False)deffn():opt.step()# Let's define a helpful benchmarking function:importtorch.utils.benchmarkasbenchmarkdefbenchmark_torch_function_in_microseconds(f,*args,**kwargs):t0=benchmark.Timer(stmt="f(*args, **kwargs)",globals={"args":args,"kwargs":kwargs,"f":f})returnt0.blocked_autorange().mean*1e6# Warmup runs to compile the functionfor_inrange(5):fn()eager_runtime=benchmark_torch_function_in_microseconds(opt.step)compiled_runtime=benchmark_torch_function_in_microseconds(fn)asserteager_runtime>compiled_runtimeprint(f"eager runtime:{eager_runtime}us")print(f"compiled runtime:{compiled_runtime}us")