diff --git a/models.py b/models.py
index 468f1ba7..1ecfdc18 100755
--- a/models.py
+++ b/models.py
@@ -150,7 +150,7 @@ class YOLOLayer(nn.Module):
             p_conf = p[..., 4]  # Conf
             p_cls = p[..., 5:]  # Class
 
-            txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG.cpu())
+            txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG)
 
             tcls = tcls[mask]
             if xy.is_cuda: