From 7d7c22cb7ea60c825fc99aa9916a229f19457108 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 26 Feb 2020 13:42:50 -0800
Subject: [PATCH] updates

---
 train.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

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: