From 84c42ab4cd8c593cc7e0436c2930164ba603526f Mon Sep 17 00:00:00 2001
From: Kevin Artica <63164026+Kahs98@users.noreply.github.com>
Date: Wed, 25 Oct 2023 16:04:43 -0500
Subject: [PATCH]  Move target tensors to GPU

---
 train.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/train.py b/train.py
index 86c7e48d5a..d88cd41e94 100644
--- a/train.py
+++ b/train.py
@@ -60,6 +60,10 @@ def train(hyp, opt, device, tb_writer=None):
     plots = not opt.evolve  # create plots
     cuda = device.type != 'cpu'
     init_seeds(2 + rank)
+    
+    # Move target tensors to GPU
+    from_which_layer = build_targets(hyp, targets, imgs)[4].to(device)
+    
     with open(opt.data) as f:
         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
     is_coco = opt.data.endswith('coco.yaml')