implement for loop

1 files changed, 6 insertions, 1 deletions

diff --git a/ot/bregman.py b/ot/bregman.py
index 6e446a1..05f7c75 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -520,7 +520,7 @@ def greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False, log=
         log['u'] = u
         log['v'] = v
 
-    while i < numItermax and stopThr_val > stopThr:
+    for i in range(numItermax):
         i += 1
         i_1 = np.argmax(np.abs(viol))
         i_2 = np.argmax(np.abs(viol_2))
@@ -547,6 +547,11 @@ def greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False, log=
 
             #print('b',np.max(abs(aviol -viol)),np.max(abs(aviol_2 - viol_2)))
 
+        if stopThr_val <= stopThr:
+            break
+    else:
+        print('Warning: Algorithm did not converge')
+
     if log:
         log['u'] = u
         log['v'] = v