def METHOD_NAME():
  """Branch function when any grad is not finite."""
  def decr_loss_scale():
    update_op = state_ops.assign(
        self._loss_scale,
        gen_math_ops.maximum(1., self._loss_scale * self._decr_ratio))
    # When loss_scale is updated, both good and bad steps are reset.
    return control_flow_ops.group(update_op, self._reset_stats())
  def just_update_steps():
    # When bad_steps is incremented, good_step is reset.
    return control_flow_ops.group(
        state_ops.assign_add(self._num_bad_steps, 1),
        state_ops.assign(self._num_good_steps, 0))
  return control_flow_ops.cond(
      self._num_bad_steps + 1 >= self._decr_every_n_nan_or_inf,
      decr_loss_scale, just_update_steps)