fix small issues

3 files changed, 7 insertions, 7 deletions

diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py
index 8cd494f..388a8d3 100644
--- a/models/pix2pix_model.py
+++ b/models/pix2pix_model.py
@@ -86,14 +86,14 @@ class Pix2PixModel(BaseModel):
     def backward_D(self):
         # Fake
         # stop backprop to the generator by detaching fake_B
-        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
+        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1).data)
         pred_fake = self.netD.forward(fake_AB.detach())
         self.loss_D_fake = self.criterionGAN(pred_fake, False)
 
         # Real
         real_AB = torch.cat((self.real_A, self.real_B), 1)
         pred_real = self.netD.forward(real_AB)
-        self.loss_D_real = self.criterionGAN(self.pred_real, True)
+        self.loss_D_real = self.criterionGAN(pred_real, True)
 
         # Combined loss
         self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
diff --git a/util/image_pool.py b/util/image_pool.py
index 5a242e6..ada1627 100644
--- a/util/image_pool.py
+++ b/util/image_pool.py
@@ -13,7 +13,7 @@ class ImagePool():
 
     def query(self, images):
         if self.pool_size == 0:
-            return images
+            return Variable(images)
         return_images = []
         for image in images:
             image = torch.unsqueeze(image, 0)
diff --git a/util/visualizer.py b/util/visualizer.py
index 22fe9da..e6e7cba 100644
--- a/util/visualizer.py
+++ b/util/visualizer.py
@@ -13,11 +13,11 @@ class Visualizer():
         self.use_html = opt.isTrain and not opt.no_html
         self.win_size = opt.display_winsize
         self.name = opt.name
+        self.opt = opt
         self.saved = False
         if self.display_id > 0:
             import visdom
             self.vis = visdom.Visdom(port=opt.display_port)
-            self.display_single_pane_ncols = opt.display_single_pane_ncols
 
         if self.use_html:
             self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
@@ -35,13 +35,13 @@ class Visualizer():
     # |visuals|: dictionary of images to display or save
     def display_current_results(self, visuals, epoch, save_result):
         if self.display_id > 0:  # show images in the browser
-            if self.display_single_pane_ncols > 0:
+            ncols = self.opt.display_single_pane_ncols
+            if ncols > 0:
                 h, w = next(iter(visuals.values())).shape[:2]
                 table_css = """<style>
                         table {border-collapse: separate; border-spacing:4px; white-space:nowrap; text-align:center}
                         table td {width: %dpx; height: %dpx; padding: 4px; outline: 4px solid black}
                         </style>""" % (w, h)
-                ncols = self.display_single_pane_ncols
                 title = self.name
                 label_html = ''
                 label_html_row = ''
@@ -76,7 +76,7 @@ class Visualizer():
                     idx += 1
 
         if self.use_html and (save_result or not self.saved):  # save images to a html file
-            self.saved  = True
+            self.saved = True
             for label, image_numpy in visuals.items():
                 img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))
                 util.save_image(image_numpy, img_path)