No CG warning when max_iter < 10
@@ -282,7 +282,11 @@ def linear_cg(
residual_norm.masked_fill_(rhs_is_zero, 0)
torch.lt(residual_norm, stop_updating_after, out=has_converged)
- if k >= 10 and bool(residual_norm.mean() < tolerance) and not (n_tridiag and k < n_tridiag_iter):
+ if (
+ k >= min(10, max_iter - 1)
+ and bool(residual_norm.mean() < tolerance)
+ and not (n_tridiag and k < min(n_tridiag_iter, max_iter - 1))
+ ):
tolerance_reached = True
break
