def get_compiled_adam_step(optimizer):
    compiled_adam = torch.compile(foreach_map_adam)
    inputs = get_inputs(optimizer)
    def compiled_adam_step():
        compiled_adam(*inputs)
    return compiled_adam_step

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