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')