text: #!/usr/bin/env python
  # -*- encoding: utf-8 -*-
  
  import torch.nn as nn
  try:
      import apex.amp as apex_amp
  except ImportError:
      pass
  
  from torch import Tensor
  
  from colossalai.nn.optimizer import ColossalaiOptimizer
  from colossalai.utils import clip_grad_norm_fp32
  
  
  class ApexAMPOptimizer(ColossalaiOptimizer):
      """ A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm
      methods
      """
  
      def backward(self, loss: Tensor):
          """Backward pass to get all gradients
  
          Args:
              loss (torch.Tensor): Loss computed by a loss function
          """
          with apex_amp.scale_loss(loss, self.optim) as scaled_loss:
              scaled_loss.backward()
  
      def clip_grad_norm(self, model: nn.Module, max_norm: float):
          """Clip gradients' norm
  
          Args:
              model (torch.nn.Module): Your model object
              max_norm (float): The max norm value for gradient clipping
          """
          if max_norm > 0:
              clip_grad_norm_fp32(apex_amp.master_params(self.optim), max_norm)
  