parameters: (self, closure=None)
return_statement: return loss
docstring: Performs a single optimization step.
  
          Args:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
docstring_summary: Performs a single optimization step.
docstring_tokens: Performs a single optimization step .
function: def step(self, closure=None):
          """Performs a single optimization step.
  
          Args:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
          """
          loss = None
          if closure is not None:
              with torch.enable_grad():
                  loss = closure()
  
          for group in self.param_groups:
              params_with_grad = []
              grads = []
              mus = []
              axs = []
              etas = []
              state_steps = []
  
              for p in group['params']:
                  if p.grad is not None:
                      params_with_grad.append(p)
                      if p.grad.is_sparse:
                          raise RuntimeError('ASGD does not support sparse gradients')
                      grads.append(p.grad)
  
                      state = self.state[p]
                      # State initialization
                      if len(state) == 0:
                          state['step'] = 0
                          state['eta'] = group['lr']
                          state['mu'] = 1
                          state['ax'] = torch.zeros_like(p, memory_format=torch.preserve_format)
  
                      mus.append(state['mu'])
                      axs.append(state['ax'])
                      etas.append(state['eta'])
  
                      state['step'] += 1
                      state_steps.append(state['step'])
  
              F.asgd(params_with_grad,
                     grads,
                     axs,
                     mus,
                     etas,
                     weight_decay=group['weight_decay'],
                     lambd=group['lambd'])
  
              # update eta and mu
              for p, mu, eta in zip(params_with_grad, mus, etas):
                  state = self.state[p]
                  state['eta'] = (group['lr'] /
                                  math.pow((1 + group['lambd'] * group['lr'] * state['step']), group['alpha']))
                  state['mu'] = 1 / max(1, state['step'] - group['t0'])
  
          return loss
function_tokens: def step ( self , closure = None ) : loss = None if closure is not None : with torch . enable_grad ( ) : loss = closure ( ) for group in self . param_groups : params_with_grad = [ ] grads = [ ] mus = [ ] axs = [ ] etas = [ ] state_steps = [ ] for p in group [ 'params' ] : if p . grad is not None : params_with_grad . append ( p ) if p . grad . is_sparse : raise RuntimeError ( 'ASGD does not support sparse gradients' ) grads . append ( p . grad ) state = self . state [ p ] # State initialization if len ( state ) == 0 : state [ 'step' ] = 0 state [ 'eta' ] = group [ 'lr' ] state [ 'mu' ] = 1 state [ 'ax' ] = torch . zeros_like ( p , memory_format = torch . preserve_format ) mus . append ( state [ 'mu' ] ) axs . append ( state [ 'ax' ] ) etas . append ( state [ 'eta' ] ) state [ 'step' ] += 1 state_steps . append ( state [ 'step' ] ) F . asgd ( params_with_grad , grads , axs , mus , etas , weight_decay = group [ 'weight_decay' ] , lambd = group [ 'lambd' ] ) # update eta and mu for p , mu , eta in zip ( params_with_grad , mus , etas ) : state = self . state [ p ] state [ 'eta' ] = ( group [ 'lr' ] / math . pow ( ( 1 + group [ 'lambd' ] * group [ 'lr' ] * state [ 'step' ] ) , group [ 'alpha' ] ) ) state [ 'mu' ] = 1 / max ( 1 , state [ 'step' ] - group [ 't0' ] ) return loss