def apply_compilation(model, optimizer):
    model = torch.compile(model, fullgraph=True, dynamic=True)
    optimizer.step = torch.compile(optimizer.step)
    return model, optimizer