From 0138dcf636c3f3f0e63110b08a8249f065e1fa73 Mon Sep 17 00:00:00 2001
From: Nathan Cassereau <84033440+ncassereau-idris@users.noreply.github.com>
Date: Mon, 1 Aug 2022 17:38:05 +0200
Subject: [MRG] Solve example throwing an error when executed on a GPU (#391)

* Solve example throwing an error when executed on a GPU

* add PR to releases.md

* update pep8 command

* pep8
---
 Makefile                                                | 2 +-
 RELEASES.md                                             | 1 +
 examples/backends/plot_sliced_wass_grad_flow_pytorch.py | 4 ++--
 ot/gromov.py                                            | 4 ++--
 ot/lp/cvx.py                                            | 2 +-
 ot/unbalanced.py                                        | 4 ++--
 6 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/Makefile b/Makefile
index 315218d..7a5cbe1 100644
--- a/Makefile
+++ b/Makefile
@@ -42,7 +42,7 @@ clean : FORCE
 	$(PYTHON) setup.py clean
 
 pep8 :
-	flake8 examples/ ot/ test/
+	flake8 examples/ ot/ test/ --count --max-line-length=127 --statistics --show-source
 
 test : FORCE pep8
 	$(PYTHON) -m pytest --durations=20 -v test/ --doctest-modules --ignore ot/gpu/  
diff --git a/RELEASES.md b/RELEASES.md
index 14d11c4..571cd74 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -19,6 +19,7 @@
 incomplete transport plan above a certain size (slightly above 46k, its square being
 roughly 2^31) (PR #381)
 - Error raised when mass mismatch in emd2 (PR #386)
+- Fixed an issue where a pytorch example would throw an error if executed on a GPU (Issue #389, PR #391)
 
 
 ## 0.8.2
diff --git a/examples/backends/plot_sliced_wass_grad_flow_pytorch.py b/examples/backends/plot_sliced_wass_grad_flow_pytorch.py
index 59e0042..f00de50 100644
--- a/examples/backends/plot_sliced_wass_grad_flow_pytorch.py
+++ b/examples/backends/plot_sliced_wass_grad_flow_pytorch.py
@@ -74,7 +74,7 @@ x_all = np.zeros((nb_iter_max, x1.shape[0], 2))
 loss_iter = []
 
 # generator for random permutations
-gen = torch.Generator()
+gen = torch.Generator(device=device)
 gen.manual_seed(42)
 
 for i in range(nb_iter_max):
@@ -136,7 +136,7 @@ x_all = np.zeros((nb_iter_max, xbary_torch.shape[0], 2))
 loss_iter = []
 
 # generator for random permutations
-gen = torch.Generator()
+gen = torch.Generator(device=device)
 gen.manual_seed(42)
 
 alpha = 0.5
diff --git a/ot/gromov.py b/ot/gromov.py
index 55ab0bd..7de8598 100644
--- a/ot/gromov.py
+++ b/ot/gromov.py
@@ -1581,7 +1581,7 @@ def gromov_barycenters(N, Cs, ps, p, lambdas, loss_fun,
 
     error = []
 
-    while(err > tol and cpt < max_iter):
+    while (err > tol and cpt < max_iter):
         Cprev = C
 
         T = [gromov_wasserstein(Cs[s], C, ps[s], p, loss_fun,
@@ -1725,7 +1725,7 @@ def fgw_barycenters(N, Ys, Cs, ps, lambdas, alpha, fixed_structure=False, fixed_
         log_['err_structure'] = []
         log_['Ts_iter'] = []
 
-    while((err_feature > tol or err_structure > tol) and cpt < max_iter):
+    while ((err_feature > tol or err_structure > tol) and cpt < max_iter):
         Cprev = C
         Xprev = X
 
diff --git a/ot/lp/cvx.py b/ot/lp/cvx.py
index fbf3c0e..361ad0f 100644
--- a/ot/lp/cvx.py
+++ b/ot/lp/cvx.py
@@ -80,7 +80,7 @@ def barycenter(A, M, weights=None, verbose=False, log=False, solver='interior-po
     if weights is None:
         weights = np.ones(A.shape[1]) / A.shape[1]
     else:
-        assert(len(weights) == A.shape[1])
+        assert len(weights) == A.shape[1]
 
     n_distributions = A.shape[1]
     n = A.shape[0]
diff --git a/ot/unbalanced.py b/ot/unbalanced.py
index 90c920c..dd9a36e 100644
--- a/ot/unbalanced.py
+++ b/ot/unbalanced.py
@@ -734,7 +734,7 @@ def barycenter_unbalanced_stabilized(A, M, reg, reg_m, weights=None, tau=1e3,
     if weights is None:
         weights = nx.ones(n_hists, type_as=A) / n_hists
     else:
-        assert(len(weights) == A.shape[1])
+        assert len(weights) == A.shape[1]
 
     if log:
         log = {'err': []}
@@ -882,7 +882,7 @@ def barycenter_unbalanced_sinkhorn(A, M, reg, reg_m, weights=None,
     if weights is None:
         weights = nx.ones(n_hists, type_as=A) / n_hists
     else:
-        assert(len(weights) == A.shape[1])
+        assert len(weights) == A.shape[1]
 
     if log:
         log = {'err': []}
-- 
cgit v1.2.3