def METHOD_NAME(self, new_scale: Optional[Union[float, torch.FloatTensor]] = None):
    """Updates the scale factor.
    If any optimizer steps were skipped, the scale is multiplied by ``backoff_factor``
    to reduce it. If ``growth_interval`` non-skipped 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 is 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 that the GradScaler uses internally.)
    .. warning::
        This method should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
        been invoked for all optimizers used this iteration.
    Args:
        new_scale (float | FloatTensor, optional):  New scale factor. (default: ``None``)
    """
    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.'
            assert isinstance(new_scale, torch.cuda.FloatTensor), reason  # type: ignore[attr-defined]
            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]
        # This is the only line changed from original grad_scaler implementation
        dist.all_reduce(found_inf_combined, reduce_operation='SUM')
        torch._amp_update_scale_(_scale, _growth_tracker, 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)