class_docstr: ## Rectified Adam Optimizer
  
  This class extends from AMSAdam optimizer defined in [`amsadam.py`](amsadam.html).
class_docstr_tok: ["#","#","Rectified","Adam","Optimizer","This","class","extends","from","AMSAdam","optimizer","defined","in","[","`","amsadam.py","`","]","(","amsadam.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$; i.e. $\sigma(.)$ and $\psi(.)$ without bias correction
      m, v = self.get_mv(state, group, grad)
  
      # Calculate $t$ the number of optimizer steps
      state["step"] += 1
  
      # Perform *RAdam* update
      self.r_adam_update(state, group, param, m, v)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/radam.py
imports: import math
  from typing import Dict, Optional
  import torch
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.amsgrad import AMSGrad
label_desc: Use image node_id 2 for calling the RAdam obj's underlying member method code with example usage: obj.step_param(state, group, grad, param) without return types