diff --git a/src/Algorithm.h b/src/Algorithm.h
index 66057cf4..1d73109c 100644
--- a/src/Algorithm.h
+++ b/src/Algorithm.h
@@ -438,7 +438,7 @@ class Algorithm {
             }
 
             // If A_U not change, U will not change and we can stop.
-            if (A_U.size() == 0 || A_U.maxCoeff() == T0 - 1) break;
+            if (this->U_size < N && (A_U.size() == 0 || A_U.maxCoeff() == T0 - 1)) break;
 
             // Update & Restore beta, A from U
             slice_restore(beta_U, U_ind, beta);