fix: GAN semantic visual output

jolibrain · Mar 22, 2024 · e7ee6bd · e7ee6bd
1 parent 5b201b8
commit e7ee6bd
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/models/base_model.py b/models/base_model.py
@@ -762,8 +762,11 @@ def get_current_visuals(self, nb_imgs, phase="train"):
                 if isinstance(name, str) and hasattr(self, name):
                     cur_visual[name] = getattr(self, name)
             visual_ret.append(cur_visual)
-            if i == nb_imgs - 1:
-                break
+            if (
+                self.opt.model_type != "cut" and self.opt.model_type != "cycle_gan"
+            ):  # GANs have more outputs in practice, including semantics
+                if i == nb_imgs - 1:
+                    break
         return visual_ret
 
     def get_display_param(self):