class_docstr: Extends SGD in PyTorch with LARS scaling from the paper
  `Large batch training of Convolutional Networks`__.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr: learning rate (default: 1e-3)
      momentum: momentum factor (default: 0)
      dampening: dampening for momentum (default: 0)
      eps: term added to the denominator to improve
          numerical stability (default: 1e-8)
      weight_decay: weight decay (L2 penalty) (default: 0)
      nesterov: enables Nesterov momentum (default: False)
      trust_coefficient: trust coefficient for computing LR (default: 0.001)
      eps: eps for division denominator (default: 1e-8)
  
  Example:
      >>> import torch_optimizer as optim
      >>> optimizer = optim.LARS(model.parameters(), lr=0.001)
      >>> optimizer.zero_grad()
      >>> loss_fn(model(input), target).backward()
      >>> optimizer.step()
  
  .. note::
      The application of momentum in the SGD part is modified according to
      the PyTorch standards. LARS scaling fits into the equation in the
      following fashion.
  
      .. math::
          \begin{aligned}
              g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\
              v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\
              p_{t+1} & = p_{t} - \text{lr} * v_{t+1},
          \\end{aligned}
  
      where :math:`p`, :math:`g`, :math:`v`, :math:`\\mu` and :math:`\beta`
      denote the parameters, gradient, velocity, momentum, and weight decay
      respectively.  The :math:`lars_lr` is defined by Eq. 6 in the paper.
      The Nesterov version is analogously modified.
  
  .. warning::
      Parameters with weight decay set to 0 will automatically be excluded
      from layer-wise LR scaling. This is to ensure consistency with papers
      like SimCLR and BYOL.
  
  
  __ https://arxiv.org/pdf/1708.03888.pdf
  
  Note:
      Reference code: https://github.com/PyTorchLightning/lightning-bolts/
docstr: Performs a single optimization step.
  
  Arguments:
      closure: A closure that reevaluates the model and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure: OptLossClosure = None) -> OptFloat:
      r"""Performs a single optimization step.
  
      Arguments:
          closure: A closure that reevaluates the model and returns the loss.
      """
      loss = None
      if closure is not None:
          with torch.enable_grad():
              loss = closure()
  
      # exclude scaling for params with 0 weight decay
      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
              p_norm = torch.norm(p.data)
              g_norm = torch.norm(p.grad.data)
  
              # lars scaling + weight decay part
              if weight_decay != 0:
                  if p_norm != 0 and g_norm != 0:
                      lars_lr = p_norm / (
                          g_norm + p_norm * weight_decay + group["eps"]
                      )
                      lars_lr *= group["trust_coefficient"]
  
                      d_p = d_p.add(p, alpha=weight_decay)
                      d_p *= lars_lr
  
              if momentum != 0:
                  param_state = self.state[p]
                  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_(d_p, alpha=1 - dampening)
                  if nesterov:
                      d_p = d_p.add(buf, alpha=momentum)
                  else:
                      d_p = buf
  
              p.add_(d_p, alpha=-group["lr"])
  
      return loss
  
imports: import torch
  from torch.optim.optimizer import Optimizer
  from .types import OptFloat, OptLossClosure, Params, State
label_desc: Use image node_id 3 for calling the LARS obj's underlying member method code with example usage: obj.step(closure) and returns: loss