code: def step(self, closure=None):
      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.float()
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adam does not support sparse gradients, please consider SparseAdam instead"
                  )
  
              p_data_fp32 = p.data.float()
  
              state = self.state[p]
  
              if len(state) == 0:
                  state["step"] = 0
                  state["exp_avg"] = torch.zeros_like(p_data_fp32)
                  state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
              else:
                  state["exp_avg"] = state["exp_avg"].type_as(
                      p_data_fp32
                  )
                  state["exp_avg_sq"] = state["exp_avg_sq"].type_as(
                      p_data_fp32
                  )
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
              beta1, beta2 = group["betas"]
  
              state["step"] += 1
  
              exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
              exp_avg.mul_(beta1).add_(1 - beta1, grad)
  
              denom = exp_avg_sq.sqrt().add_(group["eps"])
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
  
              if group["warmup"] > state["step"]:
                  scheduled_lr = (
                      1e-8
                      + state["step"] * group["lr"] / group["warmup"]
                  )
              else:
                  scheduled_lr = group["lr"]
  
              step_size = (
                  scheduled_lr
                  * math.sqrt(bias_correction2)
                  / bias_correction1
              )
  
              if group["weight_decay"] != 0:
                  p_data_fp32.add_(
                      -group["weight_decay"] * scheduled_lr, p_data_fp32
                  )
  
              p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
  
              p.data.copy_(p_data_fp32)
  
      return loss
  
file_path: jiant/jiant/ext/radam.py
imports: import math
  import torch
  from torch.optim.optimizer import Optimizer
label_desc: Use image node_id 3 for calling the AdamW obj's underlying member method code with example usage: obj.step(closure) and returns: loss