loop_example.py

import torch

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# Compile once, reuse across the training loop
compiled_model = torch.compile(model)

for batch in dataloader:
    inputs, targets = batch
    optimizer.zero_grad()
    outputs = compiled_model(inputs.cuda())
    loss = criterion(outputs, targets.cuda())
    loss.backward()
    optimizer.step()