Perform zero_grad after each batch (PythonOT#466)

eloitanguy · May 2, 2023 · 8a7035b · 8a7035b
1 parent 42a62c1
commit 8a7035b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/backends/plot_wass2_gan_torch.py b/examples/backends/plot_wass2_gan_torch.py
@@ -151,6 +151,7 @@ def forward(self, x):
 
     loss.backward()
     optimizer.step()
+    optimizer.zero_grad()
 
     del M