def METHOD_NAME(optimizer):
    for group in optimizer.param_groups:
        for p in group["params"]:
            param_state = optimizer.state[p]
            # Reinitialize momentum buffer
            if "momentum_buffer" in param_state:
                param_state["momentum_buffer"].zero_()