fix small issues

2 files changed, 5 insertions, 5 deletions

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)