original_docstring: Performs a single optimization step.
          Arguments:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
          
docstring: Performs a single optimization step.
docstring_tokens: Performs a single optimization step .
code: def step(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
code_tokens: def step ( self , closure = None ) : loss = None if closure is not None : loss = closure ( ) 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' ] 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 if self . mu != 0 : d_p . add_ ( self . mu , p . data - param_state [ 'old_init' ] ) 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 ) 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
short_docstring: Performs a single optimization step.
short_docstring_tokens: Performs a single optimization step .
comment: """Performs a single optimization step.
          Arguments:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
          """$SEP$# scale = 1**self.itr
  $SEP$# apply momentum updates
  $SEP$# apply proximal updates
  $SEP$# update accumalated local updates
  $SEP$# compute local normalizing vector a_i
  