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 _foreach_non_finite_check_and_unscale_cpu_(
      self,
      grads: List,
      found_inf: torch.Tensor,
      inv_scale: torch.Tensor,
  ) -> None:
      if len(grads) == 0:
          return
      assert (
          inv_scale.numel() == 1
      ), "inv_scale must be a 1-element tensor."
      assert (
          found_inf.numel() == 1
      ), "found_inf must be a 1-element tensor."
  
      expected_device = grads[0].device
      for tensor in grads:
          try:
              assert (
                  tensor.device == expected_device
              ), "grads must be on the same device"
          except AssertionError:
              logging.error(
                  "tensor device is %s and expected device is %s"
                  % (tensor.device, expected_device)
              )
  
          # check for non_overlapping_and_dense doesn't exist in the python world
          # as remarked here https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/AmpKernels.cu#L108
          # we assume tensor is not MTA(multi tensor apply) safe. iterate through each item regardless of dtype
          if torch.isinf(tensor).any().item() is True or torch.isnan(tensor).any().item() is True:  # type: ignore
              found_inf.data = torch.tensor([1.0])
              break
          else:
              tensor.data *= inv_scale.item()
  
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 3 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj._foreach_non_finite_check_and_unscale_cpu_(grads, found_inf, inv_scale) without return types