def apply_compilation(model, optimizer):
    model = torch.compile(model, fullgraph=True, mode="max-autotune")
    optimizer.step = torch.compile(optimizer.step)
    return model, optimizer