diff --git a/train.py b/train.py
index 5d8d3dc3..3a35d65c 100644
--- a/train.py
+++ b/train.py
@@ -254,11 +254,11 @@ def train():
             #         x['weight_decay'] = hyp['weight_decay'] * g
 
             # Plot images with bounding boxes
-            if ni == 0:
-                fname = 'train_batch%g.png' % i
-                plot_images(imgs=imgs, targets=targets, paths=paths, fname=fname)
+            if ni < 1:
+                f = 'train_batch%g.png' % i  # filename
+                plot_images(imgs=imgs, targets=targets, paths=paths, fname=f)
                 if tb_writer:
-                    tb_writer.add_image(fname, cv2.imread(fname)[:, :, ::-1], dataformats='HWC')
+                    tb_writer.add_image(f, cv2.imread(f)[:, :, ::-1], dataformats='HWC')
 
             # Multi-Scale training
             if opt.multi_scale: