class_docstr: Gradient scaler for model-parallel inf check. The inf in gradients are checked across tensor-parallel
  ranks in (1) executing optimizer step and (2) gradient scaler update.
code: def _unscale_grads_(self, optimizer, *args):
      if getattr(optimizer, "_custom_amp_unscale_grads", False):
          return optimizer.unscale_grads(*args)
      else:
          return super()._unscale_grads_(optimizer, *args)
  
file_path: apex/apex/transformer/amp/grad_scaler.py
imports: from collections import defaultdict
  import torch
  from apex.transformer import parallel_state
label_desc: Use image node_id 2 for calling the GradScaler obj's underlying member method code with example usage: obj._unscale_grads_(optimizer) and returns: optimizer, super