class_docstr: ## AdaBelief Optimizer
  
  This class extends from RAdam optimizer defined in [`radam.py`](radam.html).
class_docstr_tok: ["#","#","AdaBelief","Optimizer","This","class","extends","from","RAdam","optimizer","defined","in","[","`","radam.py","`","]","(","radam.html",")","."]
docstr: ### Take an update step for a given parameter tensor
  
  * `state` is the optimizer state of the parameter (tensor)
  * `group` stores optimizer attributes of the parameter group
  * `grad` is the current gradient tensor  $g_t$ for the parameter $      heta_{t-1}$
  * `param` is the parameter tensor $     heta_{t-1}$
docstr_tok: ["#","#","#","Take","an","update","step","for","a","given","parameter","tensor","*","`","state","`","is","the","optimizer","state","of","the","parameter","(","tensor",")","*","`","group","`","stores","optimizer","attributes","of","the","parameter","group","*","`","grad","`","is","the","current","gradient","tensor","$","g_t","$","for","the","parameter","$","heta_","{","t-1","}","$","*","`","param","`","is","the","parameter","tensor","$","heta_","{","t-1","}","$"]
code: def step_param(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      grad: torch.Tensor,
      param: torch.nn.Parameter,
  ):
      """
      ### Take an update step for a given parameter tensor
  
      * `state` is the optimizer state of the parameter (tensor)
      * `group` stores optimizer attributes of the parameter group
      * `grad` is the current gradient tensor  $g_t$ for the parameter $\theta_{t-1}$
      * `param` is the parameter tensor $\theta_{t-1}$
      """
  
      # Calculate weight decay
      grad = self.weight_decay(param, grad, group)
  
      # Get $m_t$ and $v_t$
      m, s = self.get_ms(state, group, grad)
  
      # Increment $t$ the number of optimizer steps
      state["step"] += 1
  
      if not self.rectify:
          # Perform *Adam* update, defined in [`adam.py`](adam.html), with
          # $\textcolor{cyan}{s_t} + \textcolor{red}{\epsilon}$ in place of $v_t$.
          self.adam_update(state, group, param, m, s + group["eps"])
      else:
          # Perform *Rectified Adam* update defined in [`radam.py`](radam.html), with
          # $\textcolor{cyan}{s_t} + \textcolor{red}{\epsilon}$ in place of $v_t$.
          self.r_adam_update(state, group, param, m, s + group["eps"])
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/ada_belief.py
imports: from typing import Dict, Any
  import torch
  from torch import nn
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.radam import RAdam
label_desc: Use image node_id 4 for calling the AdaBelief obj's underlying member method code with example usage: obj.step_param(state, group, grad, param) without return types