class_docstr: ShardedGradScaler helps perform gradient scaling in a shard aware manner. It extends
  functionality from GradScaler:
  * Supports Pytorch DDP and FSDP implementations
  * Support CPU offloaded tensors (as used in fully sharded data parallel[FSDP])
  * Supports the custom Mixed Precision loss dtype (fp16, bf16) that FSDP returns
  * Sync inf/nan for scaled gradient tensors on any torch.device (where tensors are placed) across
  nodes
  
  Example::
  
      # Creates a ShardedGradScaler once at the beginning of training.
      scaler = ShardedGradScaler()
  
      for epoch in epochs:
          for input, target in data:
              optimizer.zero_grad()
              output = model(input)
              loss = loss_fn(output, target)
  
              # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
              scaler.scale(loss).backward()
  
              # scaler.step() first unscales gradients of the optimizer's params.
              # If gradients don't contain infs/NaNs, optimizer.step() is then called,
              # otherwise, optimizer.step() is skipped.
              scaler.step(optimizer)
  
              # Updates the scale for next iteration.
              scaler.update()
  
  See :class:`GradScaler` for explanation of scaling/unscaling and more use cases.
  
  Args:
      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):  If ``False``, disables gradient scaling. :meth:`step` simply
          invokes the underlying ``optimizer.step()``, and other methods become no-ops.
          Default: ``True``
      process_group (ProcessGroup, optional, default=torch.distributed.group.WORLD):
          process group for sharding
docstr: If found_inf is 1.0 (True), then scale is multiplied by backoff_factor and growth_tracker is set to zero.
  Otherwise, scale is multiplied by the growth factor when the growth interval is reached.
docstr_tok: ["If","found_inf","is","1.0","(","True",")",",","then","scale","is","multiplied","by","backoff_factor","and","growth_tracker","is","set","to","zero",".","Otherwise",",","scale","is","multiplied","by","the","growth","factor","when","the","growth","interval","is","reached","."]
code: def _amp_update_scale_cpu_(self, found_inf: torch.Tensor) -> None:
      """
      If found_inf is 1.0 (True), then scale is multiplied by backoff_factor and growth_tracker is set to zero.
      Otherwise, scale is multiplied by the growth factor when the growth interval is reached.
      """
      assert (
          self._scale is not None and self._growth_tracker is not None
      )
  
      if found_inf.item() >= 1.0:
          self._scale *= self._backoff_factor
          self._growth_tracker.fill_(0)
      else:
          successful = self._growth_tracker + 1
          if successful == self._growth_interval:
              self._scale *= self._growth_factor
              self._growth_tracker.fill_(0)
          else:
              self._growth_tracker = successful
  
imports: import logging
  from collections import abc, defaultdict
  from typing import Any, Dict, Iterable, List, Optional, overload, Sequence, Tuple, Union
  import torch
  import torch.distributed
  from torch.cuda.amp.grad_scaler import _MultiDeviceReplicator, GradScaler, OptState
  from torch.distributed.distributed_c10d import ProcessGroup
label_desc: Use image node_id 10 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj._amp_update_scale_cpu_(found_inf) without return types