class_docstr: Implements Lion algorithm.
  
  Addapted from https://github.com/google/automl/tree/master/lion
  
  The Lion - EvoLved SIgn MOmeNtum - algorithm was proposed in
  https://arxiv.org/pdf/2302.06675.pdf.
  Lion aims to be more memory efficient than Adam by only tracking momentum.
  
  Caveats: As detailed in the paper, Lion requires a smaller learning rate
  lr, and larger decoupled weight decay to maintain effective weight decay
  strength. Also, the gain of Lion increases with the batch size.
  Furthermore, Lion was not found to outperform AdamW on some large language
  and text/image datasets.
  
  Arguments:
      params: iterable of parameters to optimize or dicts defining
          parameter groups
      lr: learning rate (default: 1e-3)
      betas: coefficients used for computing
          running averages of gradient and its square (default: (0.95, 0))
      weight_decay: weight decay (L2 penalty) (default: 0)
  
  Example:
      >>> import torch_optimizer as optim
      >>> optimizer = optim.Lion(model.parameters(), lr=0.001)
      >>> optimizer.zero_grad()
      >>> loss_fn(model(input), target).backward()
      >>> optimizer.step()
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()
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
  
              # Perform stepweight decay
              p.data.mul_(1 - group["lr"] * group["weight_decay"])
  
              grad = p.grad
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p)
  
              exp_avg = state["exp_avg"]
              beta1, beta2 = group["betas"]
  
              # Weight update
              update = exp_avg * beta1 + grad * (1 - beta1)
              p.add_(torch.sign(update), alpha=-group["lr"])
              # Decay the momentum running average coefficient
              exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
  
      return loss
  
imports: import torch
  from torch.optim.optimizer import Optimizer
  from .types import Betas2, OptFloat, OptLossClosure, Params
label_desc: Use image node_id 2 for calling the Lion obj's underlying member method code with example usage: obj.step(closure) and returns: loss