code: def __init__(
      self,
      params,
      lr=1e-3,
      step=0,
      bias_correction=True,
      betas=(0.9, 0.999),
      eps=1e-6,
      weight_decay=0.01,
      amsgrad=False,
      adam_w_mode=True,
      grad_averaging=True,
      max_grad_norm=1.0,
      use_nvlamb=False,
      reduced_precision_dtype=None,
  ):
      if amsgrad:
          raise RuntimeError(
              "FusedLAMB does not support the AMSGrad variant."
          )
  
      # init defaults
      defaults = dict(
          lr=torch.tensor(lr, dtype=torch.float32),
          step=torch.tensor([step], dtype=torch.int),
          bias_correction=bias_correction,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
          grad_averaging=grad_averaging,
          max_grad_norm=max_grad_norm,
      )
  
      # init base module
      super(FusedMixedPrecisionLamb, self).__init__(params, defaults)
  
      # The learning rate (lr) and optimizer step (step) should be located on device
      # in order to faciliated device sync free execution
      device = self.param_groups[0]["params"][0].device
      tensor_state = ["lr", "step"]
      for idx, group in enumerate(self.param_groups):
          for item in tensor_state:
              self.param_groups[idx][item] = group[item].to(
                  device=device
              )
  
      if multi_tensor_applier.available:
          import amp_C
  
          self.multi_tensor_l2norm = amp_C.multi_tensor_l2norm_mp
          # Skip buffer
          self._dummy_overflow_buf = torch.tensor(
              [0], dtype=torch.int, device=device
          )
          self.multi_tensor_lamb = amp_C.multi_tensor_lamb_mp
      else:
          raise RuntimeError(
              "apex.optimizers.FusedLAMB requires cuda extensions"
          )
  
      # Mixed Precision support
      self.reduced_precision_dtype = reduced_precision_dtype
      self.param_groups_full_precision = []
  
      self._step_supports_amp_scaling = True
      self.adam_w_mode = 1 if adam_w_mode else 0
      self.use_nvlamb = use_nvlamb
  
file_path: apex/apex/optimizers/fused_mixed_precision_lamb.py
imports: import torch
  from copy import deepcopy
  from itertools import chain
  from collections import defaultdict, abc
  from apex.multi_tensor_apply import multi_tensor_applier
label_desc: Use image node_id 1 to create a new FusedMixedPrecisionLamb object from inherited base classes: torch.optim with example: obj = FusedMixedPrecisionLamb(params, lr, step, bias_correction, betas, eps, weight_decay, amsgrad, adam_w_mode, grad_averaging, max_grad_norm, use_nvlamb, reduced_precision_dtype)