def METHOD_NAME(self, closure=None):
    """Performs a single optimization step.
    Arguments:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        loss = closure()
    # scale = 1**self.itr
    for group in self.param_groups:
        weight_decay = group["weight_decay"]
        momentum = group["momentum"]
        dampening = group["dampening"]
        nesterov = group["nesterov"]
        for p in group["params"]:
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            param_state = self.state[p]
            if "old_init" not in param_state:
                param_state["old_init"] = torch.clone(p.data).detach()
            local_lr = group["lr"]
            # apply momentum updates
            if momentum != 0:
                if "momentum_buffer" not in param_state:
                    buf = param_state["momentum_buffer"] = torch.clone(d_p).detach()
                else:
                    buf = param_state["momentum_buffer"]
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf
            # apply proximal updates
            if self.mu != 0:
                d_p.add_(self.mu, p.data - param_state["old_init"])
            # update accumalated local updates
            if "cum_grad" not in param_state:
                param_state["cum_grad"] = torch.clone(d_p).detach()
                param_state["cum_grad"].mul_(local_lr)
            else:
                param_state["cum_grad"].add_(local_lr, d_p)
            p.data.add_(-local_lr, d_p)
    # compute local normalizing vector a_i
    if self.momentum != 0:
        self.local_counter = self.local_counter * self.momentum + 1
        self.local_normalizing_vec += self.local_counter
    self.etamu = local_lr * self.mu
    if self.etamu != 0:
        self.local_normalizing_vec *= 1 - self.etamu
        self.local_normalizing_vec += 1
    if self.momentum == 0 and self.etamu == 0:
        self.local_normalizing_vec += 1
    self.local_steps += 1
    return loss