docstr: Performs a single optimization step.
  
  Args:
      closure (Callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["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 = []
          mu_products = []
          state_steps = []
          beta1, beta2 = group["betas"]
  
          has_complex = self._init_group(
              group,
              params_with_grad,
              grads,
              exp_avgs,
              exp_avg_sqs,
              mu_products,
              state_steps,
          )
  
          nadam(
              params_with_grad,
              grads,
              exp_avgs,
              exp_avg_sqs,
              mu_products,
              state_steps,
              beta1=beta1,
              beta2=beta2,
              lr=group["lr"],
              weight_decay=group["weight_decay"],
              momentum_decay=group["momentum_decay"],
              eps=group["eps"],
              decoupled_weight_decay=group["decoupled_weight_decay"],
              foreach=group["foreach"],
              capturable=group["capturable"],
              differentiable=group["differentiable"],
              has_complex=has_complex,
          )
  
      return loss
  
file_path: pytorch/torch/optim/nadam.py
imports: import torch
  from torch import Tensor
  from .optimizer import Optimizer, _use_grad_for_differentiable, _get_value, _dispatch_sqrt, _stack_if_compiling, _get_scalar_dtype, _default_to_fused_or_foreach, _view_as_real, _capturable_doc, _differentiable_doc, _foreach_doc
  from typing import List, Optional
label_desc: Use image node_id 4 for calling the NAdam obj's underlying member method code with example usage: obj.step(closure) and returns: loss