def METHOD_NAME(self, *args, **kwargs):
    # Copy gradients from model params to main params.
    self._set_grad_ptr()
    found_inf = self._check_overflow()
    if found_inf:
        self.optim_state = OptimState.UNSCALED    # no need to unscale grad
        self.grad_scaler.update(found_inf)    # update gradient scaler
        self._logger.info(f'Found overflow. Skip step')
        self.zero_grad()    # reset all gradients
        self._update_fp16_params()
        return
    # get combined scale. combined scale = loss scale * clipping norm
    # so that gradient = gradient / combined scale
    combined_scale = self._get_combined_scale()
    self.grad_scaler.update(found_inf)
    ret = self.optim.METHOD_NAME(div_scale=combined_scale, *args, **kwargs)
    self.zero_grad()
    self._update_fp16_params()
    return ret