diff options
| -rw-r--r-- | ot/bregman.py | 5 |
1 files changed, 1 insertions, 4 deletions
diff --git a/ot/bregman.py b/ot/bregman.py index 58e74de..c755f51 100644 --- a/ot/bregman.py +++ b/ot/bregman.py @@ -396,10 +396,7 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, log['v'] = v if nbb: # return only loss - res = np.zeros((nbb)) - for i in range(nbb): - res[i] = np.sum( - u[:, i].reshape((-1, 1)) * K * v[:, i].reshape((1, -1)) * M) + res = np.einsum('ik,ij,jk,ij->k', u, K, v, M) if log: return res, log else: |