diff options
| author | Gard Spreemann <gspr@nonempty.org> | 2023-06-14 16:51:31 +0200 |
|---|---|---|
| committer | Gard Spreemann <gspr@nonempty.org> | 2023-06-14 16:51:31 +0200 |
| commit | 96788a3fe5601e4c3f49b592aa0d9c034247862e (patch) | |
| tree | 5ee3ebcdea05f6766fc9858344913e40487e9067 /examples/backends/plot_sliced_wass_grad_flow_pytorch.py | |
| parent | 35bd2c98b642df78638d7d733bc1a89d873db1de (diff) | |
| parent | 89f1613861152432807077fbb146578611dc5888 (diff) | |
Merge tag '0.9.0' into dfsg/latestdfsg/latest
Diffstat (limited to 'examples/backends/plot_sliced_wass_grad_flow_pytorch.py')
| -rw-r--r-- | examples/backends/plot_sliced_wass_grad_flow_pytorch.py | 10 |
1 files changed, 5 insertions, 5 deletions
diff --git a/examples/backends/plot_sliced_wass_grad_flow_pytorch.py b/examples/backends/plot_sliced_wass_grad_flow_pytorch.py index cf5d64d..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): @@ -103,7 +103,7 @@ ax = pl.axis() # %% # Animate trajectories of the gradient flow along iteration -# ------------------------------------------------------- +# --------------------------------------------------------- pl.figure(3, (8, 4)) @@ -122,7 +122,7 @@ ani = animation.FuncAnimation(pl.gcf(), _update_plot, nb_iter_max, interval=100, # %% # Compute the Sliced Wasserstein Barycenter -# +# ----------------------------------------- x1_torch = torch.tensor(x1).to(device=device) x3_torch = torch.tensor(x3).to(device=device) xbinit = np.random.randn(500, 2) * 10 + 16 @@ -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 @@ -169,7 +169,7 @@ ax = pl.axis() # %% # Animate trajectories of the barycenter along gradient descent -# ------------------------------------------------------- +# ------------------------------------------------------------- pl.figure(5, (8, 4)) |