train.py GPU memory fix (#3590)

* train.py GPU memory fix

* ema

* cuda

* cuda

* zeros input

* to device

* batch index 0

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()]})