Fixing logging for small batch sizes

chenmingming · Jun 3, 2019 · 5cc5c85 · 5cc5c85
1 parent 65eac3c
commit 5cc5c85
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trainer.py b/trainer.py
@@ -540,7 +540,7 @@ def log(self, mode, inputs, outputs, losses):
         for l, v in losses.items():
             writer.add_scalar("{}".format(l), v, self.step)
 
-        for j in range(4):  # only write four images
+        for j in range(min(4, self.opt.batch_size)):  # write a maxmimum of four images
             for s in self.opt.scales:
                 for frame_id in self.opt.frame_ids:
                     writer.add_image(