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
  
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adam does not support sparse \
                                         gradients, use SparseAdam instead"
                  )
  
              is_matrix, is_need_reshape = self._check_shape(
                  grad.size()
              )
              new_shape = p.data.size()
              if is_need_reshape and group["enable_factorization"]:
                  new_shape, old_shape = self._experimental_reshape(
                      p.data.size()
                  )
                  grad = grad.view(new_shape)
  
              state = self.state[p]
              if len(state) == 0:
                  state["step"] = 0
                  if group["enable_momentum"]:
                      state["exp_avg"] = torch.zeros(
                          new_shape,
                          dtype=torch.float32,
                          device=p.grad.device,
                      )
  
                  if is_matrix and group["enable_factorization"]:
                      state["exp_avg_sq_R"] = torch.zeros(
                          (1, new_shape[1]),
                          dtype=torch.float32,
                          device=p.grad.device,
                      )
                      state["exp_avg_sq_C"] = torch.zeros(
                          (new_shape[0], 1),
                          dtype=torch.float32,
                          device=p.grad.device,
                      )
                  else:
                      state["exp_avg_sq"] = torch.zeros(
                          new_shape,
                          dtype=torch.float32,
                          device=p.grad.device,
                      )
                  if group["ams_grad"]:
                      state["exp_avg_sq_hat"] = torch.zeros(
                          new_shape,
                          dtype=torch.float32,
                          device=p.grad.device,
                      )
  
              if group["enable_momentum"]:
                  exp_avg = state["exp_avg"]
  
              if is_matrix and group["enable_factorization"]:
                  exp_avg_sq_r = state["exp_avg_sq_R"]
                  exp_avg_sq_c = state["exp_avg_sq_C"]
              else:
                  exp_avg_sq = state["exp_avg_sq"]
  
              if group["ams_grad"]:
                  exp_avg_sq_hat = state["exp_avg_sq_hat"]
  
              state["step"] += 1
              lr_t = group["lr"]
              lr_t *= max(group["eps2"], self._rms(p.data))
  
              if group["enable_momentum"]:
                  if group["non_constant_decay"]:
                      beta1_t = (
                          group["beta1"]
                          * (1 - group["beta1"] ** (state["step"] - 1))
                          / (1 - group["beta1"] ** state["step"])
                      )
                  else:
                      beta1_t = group["beta1"]
                  exp_avg.mul_(beta1_t).add_(1 - beta1_t, grad)
  
              if group["non_constant_decay"]:
                  beta2_t = (
                      group["beta2"]
                      * (1 - group["beta2"] ** (state["step"] - 1))
                      / (1 - group["beta2"] ** state["step"])
                  )
              else:
                  beta2_t = group["beta2"]
  
              if is_matrix and group["enable_factorization"]:
                  exp_avg_sq_r.mul_(beta2_t).add_(
                      1 - beta2_t,
                      torch.sum(
                          torch.mul(grad, grad).add_(group["eps1"]),
                          dim=0,
                          keepdim=True,
                      ),
                  )
                  exp_avg_sq_c.mul_(beta2_t).add_(
                      1 - beta2_t,
                      torch.sum(
                          torch.mul(grad, grad).add_(group["eps1"]),
                          dim=1,
                          keepdim=True,
                      ),
                  )
                  v = torch.mul(exp_avg_sq_c, exp_avg_sq_r).div_(
                      torch.sum(exp_avg_sq_r)
                  )
              else:
                  exp_avg_sq.mul_(beta2_t).addcmul_(
                      1 - beta2_t, grad, grad
                  ).add_((1 - beta2_t) * group["eps1"])
                  v = exp_avg_sq
  
              g = grad
              if group["enable_momentum"]:
                  g = torch.div(exp_avg, 1 - beta1_t ** state["step"])
  
              if group["ams_grad"]:
                  torch.max(exp_avg_sq_hat, v, out=exp_avg_sq_hat)
                  v = exp_avg_sq_hat
                  u = torch.div(
                      g,
                      (torch.div(v, 1 - beta2_t ** state["step"]))
                      .sqrt()
                      .add_(group["eps1"]),
                  )
              else:
                  u = torch.div(g, v.sqrt())
  
              u.div_(max(1, self._rms(u) / group["cliping_threshold"]))
              p.data.add_(
                  -lr_t
                  * (
                      u.view(old_shape)
                      if is_need_reshape
                      and group["enable_factorization"]
                      else u
                  )
              )
  
              if group["weight_decay"] != 0:
                  p.data.add_(-group["weight_decay"] * lr_t, p.data)
  
      return loss
  
imports: import torch
  import torch.optim
  from torch.nn.utils import clip_grad_norm_
  import operator
  import functools
  from copy import copy
  from math import sqrt
  import types
  import os
  import importlib
  from onmt.utils.misc import fn_args
label_desc: Use image node_id 6 for calling the AdaFactor obj's underlying member method code with example usage: obj.step(closure) and returns: loss