train.py GPU memory fix (#3590)

* train.py GPU memory fix * ema * cuda * cuda * zeros input * to device * batch index 0
ultralytics · Jun 11, 2021 · 4984cf5 · 4984cf5
1 parent e8c5237
commit 4984cf5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -335,7 +335,7 @@ def train(hyp, opt, device, tb_writer=None):
                     if tb_writer and ni == 0:
                         with warnings.catch_warnings():
                             warnings.simplefilter('ignore')  # suppress jit trace warning
-                            tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs, strict=False), [])  # graph
+                            tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
                 elif plots and ni == 10 and wandb_logger.wandb:
                     wandb_logger.log({'Mosaics': [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
                                                   save_dir.glob('train*.jpg') if x.exists()]})