From 2a163090e830e19f7ecb956e1609a2d458949037 Mon Sep 17 00:00:00 2001
From: FeepingCreature <default_357-line@yahoo.de>
Date: Mon, 6 Feb 2017 19:46:58 +0100
Subject: [PATCH] Make sure the generator training step gets a full batch of
 noise even if the data loader just ran out.

---
 main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/main.py b/main.py
index e4e5fce..8b35bb2 100644
--- a/main.py
+++ b/main.py
@@ -203,6 +203,9 @@ for epoch in range(opt.niter):
         for p in netD.parameters():
             p.requires_grad = False # to avoid computation
         netG.zero_grad()
+        # in case our last batch was the tail batch of the dataloader,
+        # make sure we feed a full batch of noise
+        noise.data.resize_(opt.batchSize, nz, 1, 1)
         noise.data.normal_(0, 1)
         fake = netG(noise)
         errG = netD(fake)