Bugfix: use nrow in correct place when visualizing results (#102)

Also normalize the result tensor when visulizaing it.
torchgan · May 1, 2019 · 375b3df · 375b3df
1 parent 9650398
commit 375b3df
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchgan/logging/visualize.py b/torchgan/logging/visualize.py
@@ -453,7 +453,7 @@ def log_console(self, trainer, image, model):
         """
         save_path = "{}/epoch{}_{}.png".format(trainer.recon, self.step, model)
         print("Generating and Saving Images to {}".format(save_path))
-        torchvision.utils.save_image(image, save_path, nrow=self.nrow)
+        torchvision.utils.save_image(image, save_path)
 
     def log_visdom(self, trainer, image, model):
         r"""Logs a generated image in visdom at the end of an epoch.
@@ -472,7 +472,7 @@ def __call__(self, trainer, **kwargs):
                 generator = getattr(trainer, model)
                 with torch.no_grad():
                     image = generator(*self.test_noise[pos])
-                    image = torchvision.utils.make_grid(image)
+                    image = torchvision.utils.make_grid(image, nrow=self.nrow, normalize=True, range=(-1, 1))
                     super(ImageVisualize, self).__call__(
                         trainer, image, model, **kwargs
                     )