updates

2019-03-04 16:11:37 +01:00
parent 5fcdcefec3
commit dc9f2ef6ba
2 changed files with 12 additions and 3 deletions
@@ -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