class_docstr: Implements NAdamW algorithm.
  
  See Table 1 in https://arxiv.org/abs/1910.05446 for the implementation of
  the NAdam algorithm (there is also a comment in the code which highlights
  the only difference of NAdamW and AdamW).
  For further details regarding the algorithm we refer to
  `Decoupled Weight Decay Regularization`_.
  
  Args:
    params (iterable): iterable of parameters to optimize or dicts defining
        parameter groups
    lr (float, optional): learning rate (default: 1e-3)
    betas (Tuple[float, float], optional): coefficients used for computing
        running averages of gradient and its square (default: (0.9, 0.999))
    eps (float, optional): term added to the denominator to improve
        numerical stability (default: 1e-8)
    weight_decay (float, optional): weight decay coefficient (default: 1e-2)
  .. _Decoupled Weight Decay Regularization:
      https://arxiv.org/abs/1711.05101
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
docstr: Performs a single optimization step.
  
  Args:
    closure (callable, optional): A closure that reevaluates the model
        and returns the loss.
code: def step(self, closure=None):
      """Performs a single optimization step.
  
      Args:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
      """
      self._cuda_graph_capture_health_check()
  
      loss = None
      if closure is not None:
          with torch.enable_grad():
              loss = closure()
  
      for group in self.param_groups:
          params_with_grad = []
          grads = []
          exp_avgs = []
          exp_avg_sqs = []
          state_steps = []
          beta1, beta2 = group["betas"]
  
          for p in group["params"]:
              if p.grad is None:
                  continue
              params_with_grad.append(p)
              if p.grad.is_sparse:
                  raise RuntimeError(
                      "NAdamW does not support sparse gradients"
                  )
              grads.append(p.grad)
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = torch.tensor(0.0)
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(
                      p, memory_format=torch.preserve_format
                  )
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(
                      p, memory_format=torch.preserve_format
                  )
  
              exp_avgs.append(state["exp_avg"])
              exp_avg_sqs.append(state["exp_avg_sq"])
              state_steps.append(state["step"])
  
          nadamw(
              params_with_grad,
              grads,
              exp_avgs,
              exp_avg_sqs,
              state_steps,
              beta1=beta1,
              beta2=beta2,
              lr=group["lr"],
              weight_decay=group["weight_decay"],
              eps=group["eps"],
              maximize=group["maximize"],
              capturable=group["capturable"],
          )
  
      return loss
  
file_path: pytorch-image-models/timm/optim/nadamw.py
imports: import math
  from typing import List, Optional
  import torch
  from torch import Tensor
label_desc: Use image node_id 3 for calling the NAdamW obj's underlying member method code with example usage: obj.step(closure) and returns: loss