class_docstr: Optimizer wrapper for mixed precision training in FP16 using PyTorch AMP.
  
  Args:
      optim (Optimizer): Optimizer to wrap.
      init_scale (float): Initial scale factor. Default: 2**16.
      growth_factor (float): Factor by which the scale is multiplied during
          :meth:`torch.cuda.amp.GradScaler.step` if gradients were found to be finite
          this iteration. Default: 2.0.
      backoff_factor (float): Factor by which the scale is multiplied during
          :meth:`torch.cuda.amp.GradScaler.step` if gradients were found to be infinite
          this iteration. Default: 0.5.
      growth_interval (int): Number of iterations between :meth:`torch.cuda.amp.GradScaler.step`
          calls that may cause the scale to increase. Default: 2000.
class_docstr_tok: ["Optimizer","wrapper","for","mixed","precision","training","in","FP16","using","PyTorch","AMP",".","Args",":","optim","(","Optimizer",")",":","Optimizer","to","wrap",".","init_scale","(","float",")",":","Initial","scale","factor",".","Default",":","2","*","*","16.","growth_factor","(","float",")",":","Factor","by","which","the","scale","is","multiplied","during",":","meth",":","`","torch.cuda.amp.GradScaler.step","`","if","gradients","were","found","to","be","finite","this","iteration",".","Default",":","2.0.","backoff_factor","(","float",")",":","Factor","by","which","the","scale","is","multiplied","during",":","meth",":","`","torch.cuda.amp.GradScaler.step","`","if","gradients","were","found","to","be","infinite","this","iteration",".","Default",":","0.5.","growth_interval","(","int",")",":","Number","of","iterations","between",":","meth",":","`","torch.cuda.amp.GradScaler.step","`","calls","that","may","cause","the","scale","to","increase",".","Default",":","2000","."]
code: def step(self, *args, **kwargs) -> Optional[float]:
      out = self.scaler.step(self.optim, *args, **kwargs)
      self.scaler.update()
      return out
  
imports: from typing import Callable, Optional, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
  from torch.optim import Optimizer
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from colossalai.utils.device import autocast
  from .mixed_precision_base import MixedPrecision
label_desc: Use image node_id 3 for calling the TorchAMPOptimizer obj's underlying member method code with example usage: obj.step() and returns: out