class_docstr: Override from Pytorch implementation to ensure aggregations done in FP32.
class_docstr_tok: ["Override","from","Pytorch","implementation","to","ensure","aggregations","done","in","FP32","."]
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()  # NOTE: cast to FP32 here
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adam does not support sparse gradients, "
                      "please consider SparseAdam instead"
                  )
              amsgrad = group["amsgrad"]
  
              p_data_fp32 = p.data.float()  # NOTE: cast to FP32 here
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p_data_fp32)
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
                  if amsgrad:
                      # Maintains max of all exp. moving avg. of sq. grad. values
                      state["max_exp_avg_sq"] = torch.zeros_like(
                          p_data_fp32
                      )
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
              if amsgrad:
                  max_exp_avg_sq = state["max_exp_avg_sq"]
              beta1, beta2 = group["betas"]
  
              state["step"] += 1
  
              if group["weight_decay"] != 0:
                  grad.add_(p_data_fp32, alpha=group["weight_decay"])
  
              # Decay the first and second moment running average coefficient
              exp_avg.mul_(beta1).add_(grad, alpha=(1 - beta1))
              exp_avg_sq.mul_(beta2).addcmul_(
                  grad, grad, value=(1 - beta2)
              )
              if amsgrad:
                  # Maintains the maximum of all 2nd moment running avg. till now
                  torch.max(
                      max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq
                  )
                  # Use the max. for normalizing running avg. of gradient
                  denom = max_exp_avg_sq.sqrt().add_(group["eps"])
              else:
                  denom = exp_avg_sq.sqrt().add_(group["eps"])
  
              bias_correction1 = 1 - beta1 ** state["step"]
              bias_correction2 = 1 - beta2 ** state["step"]
              step_size = (
                  group["lr"]
                  * math.sqrt(bias_correction2)
                  / bias_correction1
              )
  
              p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size)
  
              p.data.copy_(p_data_fp32)
  
      return loss
  
imports: import math
  from itertools import chain
  from typing import Optional
  import parlai.utils.logging
label_desc: Use image node_id 1 for calling the MemoryEfficientFP16Adam obj's underlying member method code with example usage: obj.step(closure) and returns: loss