class_docstr: `AdaBelief Optimizer, adapting stepsizes by the belief in observed gradients`
  Paper: https://arxiv.org/abs/2010.07468
  Implementation has been copied over from the original author (https://github.com/juntang-zhuang/Adabelief-Optimizer)
class_docstr_tok: ["`","AdaBelief","Optimizer",",","adapting","stepsizes","by","the","belief","in","observed","gradients","`","Paper",":","https",":","\/\/arxiv.org\/abs\/2010.07468","Implementation","has","been","copied","over","from","the","original","author","(","https",":","\/\/github.com\/juntang-zhuang\/Adabelief-Optimizer",")"]
docstr: Performs a single optimization step.
  
  Arguments:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None, **kwargs):
      """Performs a single optimization step.
  
      Arguments:
          closure (callable, optional): 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:
          for p in group["params"]:
              if p.grad is None:
                  continue
              grad = p.grad.data
              if grad.is_sparse:
                  raise RuntimeError(
                      "AdaBelief does not support sparse gradients, please consider SparseAdam instead"
                  )
              amsgrad = group["amsgrad"]
  
              state = self.state[p]
  
              beta1, beta2 = group["betas"]
  
              # State initialization
              if len(state) == 0:
                  state["rho_inf"] = 2.0 / (1.0 - beta2) - 1.0
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(
                      p.data, memory_format=torch.preserve_format
                  )
                  # Exponential moving average of squared gradient values
                  state["exp_avg_var"] = torch.zeros_like(
                      p.data, memory_format=torch.preserve_format
                  )
                  if amsgrad:
                      # Maintains max of all exp. moving avg. of sq. grad. values
                      state["max_exp_avg_var"] = torch.zeros_like(
                          p.data, memory_format=torch.preserve_format
                      )
  
              # get current state variable
              exp_avg, exp_avg_var = (
                  state["exp_avg"],
                  state["exp_avg_var"],
              )
  
              state["step"] += 1
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
  
              # perform weight decay, check if decoupled weight decay
              if self.weight_decouple:
                  if not self.fixed_decay:
                      p.data.mul_(
                          1.0 - group["lr"] * group["weight_decay"]
                      )
                  else:
                      p.data.mul_(1.0 - group["weight_decay"])
              else:
                  if group["weight_decay"] != 0:
                      grad.add_(group["weight_decay"], p.data)
  
              # Update first and second moment running average
              exp_avg.mul_(beta1).add_(1 - beta1, grad)
              grad_residual = grad - exp_avg
              exp_avg_var.mul_(beta2).addcmul_(
                  1 - beta2, grad_residual, grad_residual
              )
  
              if amsgrad:
                  max_exp_avg_var = state["max_exp_avg_var"]
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch.max(
                      max_exp_avg_var, exp_avg_var, out=max_exp_avg_var
                  )
  
                  # Use the max. for normalizing running avg. of gradient
                  denom = (
                      max_exp_avg_var.add_(group["eps"]).sqrt()
                      / math.sqrt(bias_correction2)
                  ).add_(group["eps"])
              else:
                  denom = (
                      exp_avg_var.add_(group["eps"]).sqrt()
                      / math.sqrt(bias_correction2)
                  ).add_(group["eps"])
  
              if not self.rectify:
                  # Default update
                  step_size = group["lr"] / bias_correction1
                  p.data.addcdiv_(-step_size, exp_avg, denom)
  
              else:  # Rectified update
                  # calculate rho_t
                  state["rho_t"] = state["rho_inf"] - 2 * state[
                      "step"
                  ] * beta2 ** state["step"] / (
                      1.0 - beta2 ** state["step"]
                  )
  
                  if (
                      state["rho_t"] > 4
                  ):  # perform Adam style update if variance is small
                      rho_inf, rho_t = state["rho_inf"], state["rho_t"]
                      rt = (
                          (rho_t - 4.0)
                          * (rho_t - 2.0)
                          * rho_inf
                          / (rho_inf - 4.0)
                          / (rho_inf - 2.0)
                          / rho_t
                      )
                      rt = math.sqrt(rt)
  
                      step_size = rt * group["lr"] / bias_correction1
  
                      p.data.addcdiv_(-step_size, exp_avg, denom)
  
                  else:  # perform SGD style update
                      p.data.add_(-group["lr"], exp_avg)
  
      return loss
  
file_path: pytext/pytext/optimizer/adabelief.py
imports: import math
  import torch
  from torch.optim.optimizer import Optimizer
  from .optimizers import Optimizer
label_desc: Use image node_id 5 for calling the AdaBelief obj's underlying member method code with example usage: obj.step(closure) and returns: loss