class_docstr: Implements Aggregated Momentum Gradient Descent.
  
  It has been proposed in `Aggregated Momentum: Stability Through Passive
  Damping`__
  
  Example:
      >>> import torch_optimizer as optim
      >>> optimizer = optim.AggMo(model.parameters(), lr=0.1)
      >>> optimizer.zero_grad()
      >>> loss_fn(model(input), target).backward()
      >>> optimizer.step()
  
   __ https://arxiv.org/abs/1804.00325
  
  Note:
      Reference code: https://github.com/AtheMathmo/AggMo/blob/master/aggmo.py  # noqa
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:
          loss = closure()
  
      for group in self.param_groups:
          weight_decay = group["weight_decay"]
          betas = group["betas"]
          total_mom = float(len(betas))
  
          for p in group["params"]:
              if p.grad is None:
                  continue
              d_p = p.grad.data
              if weight_decay != 0:
                  d_p.add_(p.data, alpha=weight_decay)
              param_state = self.state[p]
              if "momentum_buffer" not in param_state:
                  param_state["momentum_buffer"] = {}
                  for beta in betas:
                      param_state["momentum_buffer"][
                          beta
                      ] = torch.zeros_like(
                          p.data, memory_format=torch.preserve_format
                      )
              for beta in betas:
                  buf = param_state["momentum_buffer"][beta]
                  buf.mul_(beta).add_(d_p)
                  p.data.sub_(buf, alpha=group["lr"] / total_mom)
      return loss
  
file_path: pytorch-optimizer/torch_optimizer/aggmo.py
imports: from typing import List, Tuple, Type, TypeVar, Union
  import torch
  from torch.optim.optimizer import Optimizer
  from .types import OptFloat, OptLossClosure, Params
label_desc: Use image node_id 3 for calling the AggMo obj's underlying member method code with example usage: obj.step(closure) and returns: loss