From 256a3e89d2d78ab90eff44a4771de2586f62b15c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 25 Jun 2020 17:52:56 -0700 Subject: [PATCH] small dataset bug fix #140 --- train.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index ea72f15..94139a0 100644 --- a/train.py +++ b/train.py @@ -287,10 +287,10 @@ def train(hyp): # Plot if ni < 3: - f = 'train_batch%g.jpg' % i # filename - res = plot_images(images=imgs, targets=targets, paths=paths, fname=f) - if tb_writer: - tb_writer.add_image(f, res, dataformats='HWC', global_step=epoch) + f = 'train_batch%g.jpg' % ni # filename + result = plot_images(images=imgs, targets=targets, paths=paths, fname=f) + if tb_writer and result is not None: + tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) # tb_writer.add_graph(model, imgs) # add model to tensorboard # end batch ------------------------------------------------------------------------------------------------