class_docstr: A shard aware Grad Scaler which enables loss scaling with/without cpu_offload. This is a
  slight modification of the pytorch grad scaler.
  https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
class_docstr_tok: ["A","shard","aware","Grad","Scaler","which","enables","loss","scaling","with\/without","cpu_offload",".","This","is","a","slight","modification","of","the","pytorch","grad","scaler",".","https",":","\/\/pytorch.org\/docs\/stable\/amp.html","#","torch.cuda.amp.GradScaler"]
code: def unscale_(self, optimizer: SGD) -> None:  # type: ignore
      if not self._enabled:
          return
  
      super()._check_scale_growth_tracker("unscale_")  # type: ignore
  
      optimizer_state = self._per_optimizer_states[id(optimizer)]
  
      if optimizer_state["stage"] is OptState.UNSCALED:
          raise RuntimeError(
              "unscale_() has already been called on this optimizer since the last update()."
          )
      elif optimizer_state["stage"] is OptState.STEPPED:
          raise RuntimeError("unscale_() is being called after step().")
  
      # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
      assert self._scale is not None
      inv_scale = self._scale.double().reciprocal().float()
      found_inf = torch.full(
          (1,), 0.0, dtype=torch.float32, device=self._scale.device
      )
  
      optimizer_state["found_inf_per_device"] = self._unscale_grads_(
          optimizer, inv_scale, found_inf, True
      )
      optimizer_state["stage"] = OptState.UNSCALED
  
      # Synchronize the detected inf across the ranks
      optimizer_state = self._per_optimizer_states[id(optimizer)]
      last_handle = None
  
      for v in optimizer_state["found_inf_per_device"].values():
          if v.device.type == "cpu":
              v_on_cuda = v.cuda()
              last_handle = dist.all_reduce(
                  v_on_cuda, async_op=True, group=self.group
              )
              v_on_cuda.cpu()
          else:
              last_handle = dist.all_reduce(
                  v, async_op=True, group=self.group
              )
  
      # Make sure that the calls are done before moving out.
      # The calls are executed in sequence, waiting for the last one is enough
      if last_handle is not None:
          last_handle.wait()
  
imports: from collections import defaultdict
  import collections.abc
  from enum import Enum
  import logging
  from typing import Any, Dict, List, Optional, Union
  import warnings
  import torch
  from torch.cuda import FloatTensor
  from torch.cuda.amp.common import amp_definitely_not_available
  from torch.cuda.amp.grad_scaler import GradScaler
  import torch.distributed
  from torch.optim import Optimizer
  from torch.optim.sgd import SGD
  from fairscale.internal import torch_version
label_desc: Use image node_id 5 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj.unscale_(optimizer) without return types