From 5e0cce771eb26ac85c44b8283835d17ea6c183ee Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 4 Mar 2019 15:52:51 +0100 Subject: [PATCH] updates --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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: