def METHOD_NAME(self):
    """
    Update the scale factor and clear information about infinities.
    This method should be called after each optimization step.
    """
    if not self._enabled:
        return
    found_inf_device = torch.full(
        (1,), 1.0 if self._found_inf else 0.0, device=self._scale.device
    )
    torch._amp_update_scale_(
        self._scale,
        self._growth_tracker,
        found_inf_device,
        self._growth_factor,
        self._backoff_factor,
        self._growth_interval,
    )
    # Clear infinity found status
    self._found_inf = False