def METHOD_NAME(self, opt):
    opt.local_counter = 0
    opt.local_normalizing_vec = 0
    opt.local_steps = 0
    for group in opt.param_groups:
        for p in group["params"]:
            param_state = opt.state[p]
            param_state["cum_grad"].zero_()
            # Reinitialize momentum buffer
            if "momentum_buffer" in param_state:
                param_state["momentum_buffer"].zero_()