Perform zero_grad after each batch (#466)

Author: kachayev

examples/backends/plot_wass2_gan_torch.py

@@ -151,6 +151,7 @@ def forward(self, x):
 
     loss.backward()
     optimizer.step()
+    optimizer.zero_grad()
 
     del M