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