From 1a740d5b3bdada9ab4173bddb0b808e735b7da4c Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 11 Jun 2021 20:24:03 +0200
Subject: [PATCH] train.py GPU memory fix (#3590)

* train.py GPU memory fix

* ema

* cuda

* cuda

* zeros input

* to device

* batch index 0

(cherry picked from commit 4984cf54be4eb88f00ccf33a05f57681b2a770ab)
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 29a0ab36a421..90470b66ea3c 100644
--- a/train.py
+++ b/train.py
@@ -371,7 +371,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()]})