glenn-jocher commited on
Commit
b7cd1f5
·
unverified ·
1 Parent(s): 3f74cd9

TensorBoard add_graph() fix (#3236)

Browse files
Files changed (1) hide show
  1. train.py +3 -3
train.py CHANGED
@@ -330,9 +330,9 @@ def train(hyp, opt, device, tb_writer=None):
330
  if plots and ni < 3:
331
  f = save_dir / f'train_batch{ni}.jpg' # filename
332
  Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
333
- # if tb_writer:
334
- # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
335
- # tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
336
  elif plots and ni == 10 and wandb_logger.wandb:
337
  wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
338
  save_dir.glob('train*.jpg') if x.exists()]})
 
330
  if plots and ni < 3:
331
  f = save_dir / f'train_batch{ni}.jpg' # filename
332
  Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
333
+ if tb_writer:
334
+ tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
335
+ # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
336
  elif plots and ni == 10 and wandb_logger.wandb:
337
  wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
338
  save_dir.glob('train*.jpg') if x.exists()]})