docstr: Updates the scale factor.
  
  If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
  to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
  the scale is multiplied by ``growth_factor`` to increase it.
  
  Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
  used directly, it's used to fill GradScaler's internal scale tensor. So if
  ``new_scale`` was a tensor, later in-place changes to that tensor will not further
  affect the scale GradScaler uses internally.)
  
  Args:
      new_scale (float or :class:`torch.cuda.FloatTensor`, optional, default=None):  New scale factor.
  
  .. warning::
      :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
      been invoked for all optimizers used this iteration.
code: def update(self, new_scale=None):
      """
      Updates the scale factor.
  
      If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
      to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
      the scale is multiplied by ``growth_factor`` to increase it.
  
      Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
      used directly, it's used to fill GradScaler's internal scale tensor. So if
      ``new_scale`` was a tensor, later in-place changes to that tensor will not further
      affect the scale GradScaler uses internally.)
  
      Args:
          new_scale (float or :class:`torch.cuda.FloatTensor`, optional, default=None):  New scale factor.
  
      .. warning::
          :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
          been invoked for all optimizers used this iteration.
      """
      if not self._enabled:
          return
  
      _scale, _growth_tracker = self._check_scale_growth_tracker(
          "update"
      )
  
      if new_scale is not None:
          # Accept a new user-defined scale.
          if isinstance(new_scale, float):
              self._scale.fill_(new_scale)  # type: ignore[union-attr]
          else:
              reason = "new_scale should be a float or a 1-element torch.cuda.FloatTensor with requires_grad=False."
              # type: ignore[attr-defined]
              assert isinstance(
                  new_scale, torch.cuda.FloatTensor
              ), reason
              assert new_scale.numel() == 1, reason
              assert new_scale.requires_grad is False, reason
              self._scale.copy_(new_scale)  # type: ignore[union-attr]
      else:
          # Consume shared inf/nan data collected from optimizers to update the scale.
          # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous.
          found_infs = [
              found_inf.to(device=_scale.device, non_blocking=True)
              for state in self._per_optimizer_states.values()
              for found_inf in state["found_inf_per_device"].values()
          ]
  
          assert (
              len(found_infs) > 0
          ), "No inf checks were recorded prior to update."
  
          found_inf_combined = found_infs[0]
          if len(found_infs) > 1:
              for i in range(1, len(found_infs)):
                  found_inf_combined += found_infs[i]
  
          if self._higher_than_torch18:
              torch._amp_update_scale_(
                  _scale,
                  _growth_tracker,
                  found_inf_combined,
                  self._growth_factor,
                  self._backoff_factor,
                  self._growth_interval,
              )
          else:
              self._scale = torch._amp_update_scale(
                  _growth_tracker,
                  _scale,
                  found_inf_combined,
                  self._growth_factor,
                  self._backoff_factor,
                  self._growth_interval,
              )
  
      # To prepare for next iteration, clear the data collected from optimizers this iteration.
      self._per_optimizer_states = defaultdict(
          _refresh_per_optimizer_state
      )
  
imports: import warnings
  from collections import abc, defaultdict
  from enum import Enum
  from typing import Any, Dict, List, Optional, Tuple
  import torch
  import torch.distributed
  from packaging import version
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
label_desc: Use image node_id 9 for calling the GradScaler obj's underlying member method code with example usage: obj.update(new_scale) without return types