updates

models.py

@@ -150,8 +150,6 @@ class YOLOLayer(nn.Module):
             p_conf = p[..., 4]  # Conf
             p_cls = p[..., 5:]  # Class
 
-            print(self.anchor_vec.device, self.anchor_vec, self.nA, self.nC, nG)
-
             txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG)
 
             tcls = tcls[mask]

utils/utils.py

@@ -256,6 +256,9 @@ def build_targets(target, anchor_vec, nA, nC, nG):
         # iou of targets-anchors (using wh only)
         box1 = gwh
         box2 = anchor_vec.unsqueeze(1)
+
+        print(box1.device,box2.device)
+
         inter_area = torch.min(box1, box2).prod(2)
         iou = inter_area / (box1.prod(1) + box2.prod(2) - inter_area + 1e-16)