class_docstr: Implements the AdaHessian algorithm from "ADAHESSIAN: An Adaptive Second OrderOptimizer for Machine Learning"
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining parameter groups
      lr (float, optional): learning rate (default: 0.1)
      betas ((float, float), optional): coefficients used for computing running averages of gradient and the
          squared hessian trace (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 (L2 penalty) (default: 0.0)
      hessian_power (float, optional): exponent of the hessian trace (default: 1.0)
      update_each (int, optional): compute the hessian trace approximation only after *this* number of steps
          (to save time) (default: 1)
      n_samples (int, optional): how many times to sample `z` for the approximation of the hessian trace (default: 1)
docstr: Performs a single optimization step.
  Arguments:
      closure (callable, optional) -- a closure that reevaluates the model and returns the loss (default: None)
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")","--","a","closure","that","reevaluates","the","model","and","returns","the","loss","(","default",":","None",")"]
code: def step(self, closure=None):
      """
      Performs a single optimization step.
      Arguments:
          closure (callable, optional) -- a closure that reevaluates the model and returns the loss (default: None)
      """
  
      loss = None
      if closure is not None:
          loss = closure()
  
      self.zero_hessian()
      self.set_hessian()
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None or p.hess is None:
                  continue
  
              if self.avg_conv_kernel and p.dim() == 4:
                  p.hess = (
                      torch.abs(p.hess)
                      .mean(dim=[2, 3], keepdim=True)
                      .expand_as(p.hess)
                      .clone()
                  )
  
              # Perform correct stepweight decay as in AdamW
              p.mul_(1 - group["lr"] * group["weight_decay"])
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 1:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p)
                  # Exponential moving average of Hessian diagonal square values
                  state["exp_hessian_diag_sq"] = torch.zeros_like(p)
  
              exp_avg, exp_hessian_diag_sq = (
                  state["exp_avg"],
                  state["exp_hessian_diag_sq"],
              )
              beta1, beta2 = group["betas"]
              state["step"] += 1
  
              # Decay the first and second moment running average coefficient
              exp_avg.mul_(beta1).add_(p.grad, alpha=1 - beta1)
              exp_hessian_diag_sq.mul_(beta2).addcmul_(
                  p.hess, p.hess, value=1 - beta2
              )
  
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
  
              k = group["hessian_power"]
              denom = (
                  (exp_hessian_diag_sq / bias_correction2)
                  .pow_(k / 2)
                  .add_(group["eps"])
              )
  
              # make update
              step_size = group["lr"] / bias_correction1
              p.addcdiv_(exp_avg, denom, value=-step_size)
  
      return loss
  
file_path: pywick/pywick/optimizers/adahessian.py
imports: import torch
label_desc: Use image node_id 6 for calling the Adahessian obj's underlying member method code with example usage: obj.step(closure) and returns: loss