class_docstr: A wrapper class which integrate Pytorch AMP with an optimizer
  
  Args:
      optim (torch.optim.Optimizer): A normal optimizer like Adam or SGD.
      init_scale (float, optional, default=2.**16):  Initial scale factor.
      growth_factor (float, optional, default=2.0):  Factor by which the scale is multiplied during
          :meth:`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.
      backoff_factor (float, optional, default=0.5):  Factor by which the scale is multiplied during
          :meth:`update` if inf/NaN gradients occur in an iteration.
      growth_interval (int, optional, default=2000):  Number of consecutive iterations without inf/NaN gradients
          that must occur for the scale to be multiplied by ``growth_factor``.
      enabled (bool, optional, default=True):  If ``False``, disables gradient scaling. :meth:`step` simply
          invokes the underlying ``optimizer.step()``, and other methods become no-ops.
code: def __init__(self, optim: Optimizer, *args, **kwargs):
      super().__init__(optim)
      self.scaler = GradScaler(*args, **kwargs)
  
imports: from colossalai.utils.device import autocast
  import torch.nn
  from torch import Tensor
  from torch.nn.modules.loss import _Loss
  from torch.optim import Optimizer
  from colossalai.interface import OptimizerWrapper
  from colossalai.legacy.utils import clip_grad_norm_fp32
  from ._grad_scaler import GradScaler
label_desc: Use image node_id 1 to create a new TorchAMPOptimizer object from inherited base classes: OptimizerWrapper with example: obj = TorchAMPOptimizer(optim)