From a919f960df57b4ea4beff57a3f7262b8064d8159 Mon Sep 17 00:00:00 2001
From: Hicham Janati <hicham.janati@inria.fr>
Date: Sat, 19 Oct 2019 12:05:59 +0200
Subject: same for unbalanced

---
 ot/unbalanced.py | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'ot/unbalanced.py')

diff --git a/ot/unbalanced.py b/ot/unbalanced.py
index d516dfc..978df08 100644
--- a/ot/unbalanced.py
+++ b/ot/unbalanced.py
@@ -1002,12 +1002,14 @@ def barycenter_unbalanced(A, M, reg, reg_m, method="sinkhorn", weights=None,
 
     if method.lower() == 'sinkhorn':
         return barycenter_unbalanced_sinkhorn(A, M, reg, reg_m,
+                                              weights=weights,
                                               numItermax=numItermax,
                                               stopThr=stopThr, verbose=verbose,
                                               log=log, **kwargs)
 
     elif method.lower() == 'sinkhorn_stabilized':
         return barycenter_unbalanced_stabilized(A, M, reg, reg_m,
+                                                weights=weights,
                                                 numItermax=numItermax,
                                                 stopThr=stopThr,
                                                 verbose=verbose,
@@ -1015,6 +1017,7 @@ def barycenter_unbalanced(A, M, reg, reg_m, method="sinkhorn", weights=None,
     elif method.lower() in ['sinkhorn_reg_scaling']:
         warnings.warn('Method not implemented yet. Using classic Sinkhorn Knopp')
         return barycenter_unbalanced(A, M, reg, reg_m,
+                                     weights=weights,
                                      numItermax=numItermax,
                                      stopThr=stopThr, verbose=verbose,
                                      log=log, **kwargs)
-- 
cgit v1.2.3