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()
    for group in self.param_groups:
        for p in group['params']:
            if p.grad is None:
                continue
            grad = p.grad.data
            if grad.is_sparse:
                raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
            state = self.state[p]
            # State initialization
            if not state:
                state['step'] = 0
                # Exponential moving average of gradient values
                state['next_m'] = torch.zeros_like(p.data)
                # Exponential moving average of squared gradient values
                state['next_v'] = torch.zeros_like(p.data)
            next_m, next_v = state['next_m'], state['next_v']
            beta1, beta2 = group['b1'], group['b2']
            # Add grad clipping
            if group['max_grad_norm'] > 0:
                clip_grad_norm_(p, group['max_grad_norm'])
            # Decay the first and second moment running average coefficient
            # In-place operations to update the averages at the same time
            next_m.mul_(beta1).add_(1 - beta1, grad)
            next_v.mul_(beta2).addcmul_(1 - beta2, grad, grad)
            update = next_m / (next_v.sqrt() + group['e'])
            # Just adding the square of the weights to the loss function is *not*
            # the correct way of using L2 regularization/weight decay with Adam,
            # since that will interact with the m and v parameters in strange ways.
            #
            # Instead we want to decay the weights in a manner that doesn't interact
            # with the m/v parameters. This is equivalent to adding the square
            # of the weights to the loss with plain (non-momentum) SGD.
            if group['weight_decay_rate'] > 0.0:
                update += group['weight_decay_rate'] * p.data
            if group['t_total'] != -1:
                schedule_fct = SCHEDULES[group['schedule']]
                lr_scheduled = group['lr'] * schedule_fct(state['step']/group['t_total'], group['warmup'])
            else:
                lr_scheduled = group['lr']
            update_with_lr = lr_scheduled * update
            p.data.add_(-update_with_lr)
            state['step'] += 1
            # step_size = lr_scheduled * math.sqrt(bias_correction2) / bias_correction1
            # No bias correction
            # bias_correction1 = 1 - beta1 ** state['step']
            # bias_correction2 = 1 - beta2 ** state['step']
    return loss