diff --git a/train.py b/train.py index 9aff01a42c08..f0ad206d4b94 100644 --- a/train.py +++ b/train.py @@ -318,8 +318,9 @@ def train(hyp, opt, device, tb_writer=None, wandb=None): # if tb_writer: # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) # tb_writer.add_graph(model, imgs) # add model to tensorboard - elif plots and ni == 3 and wandb: - wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')]}) + elif plots and ni > 3 and wandb: + wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg') + if x.exists()]}) # end batch ------------------------------------------------------------------------------------------------ # end epoch ----------------------------------------------------------------------------------------------------