diff --git a/models.py b/models.py index f3d57e68..8f7f79d6 100755 --- a/models.py +++ b/models.py @@ -150,10 +150,13 @@ 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) + if p.is_cuda: + txy, twh, mask, tcls = build_targets(targets, self.anchor_vec.cuda(), self.nA, self.nC, nG) + else: + txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG) tcls = tcls[mask] - if xy.is_cuda: + if p.is_cuda: txy, twh, mask, tcls = txy.cuda(), twh.cuda(), mask.cuda(), tcls.cuda() # Compute losses diff --git a/utils/utils.py b/utils/utils.py index 2e45ddb4..e78a3e62 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -242,8 +242,14 @@ def build_targets(target, anchor_vec, nA, nC, nG): tconf = torch.ByteTensor(nB, nA, nG, nG).fill_(0) tcls = torch.ByteTensor(nB, nA, nG, nG, nC).fill_(0) # nC = number of classes + if anchor_vec.is_cuda(): + txy = txy.cuda() + twh = twh.cuda() + tconf = tconf.cuda() + tcls = tcls.cuda() + for b in range(nB): - t = target[b].cpu() + t = target[b] nTb = len(t) # number of targets if nTb == 0: continue