minor fix

FL77N · Jul 21, 2022 · 2596994 · 2596994
1 parent 894a93e
commit 2596994
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/detect.py b/detect.py
@@ -73,7 +73,7 @@ def detect(save_img=False):
 
         # Inference
         t1 = time_synchronized()
-        pred = model(img, augment=opt.augment)
+        pred = model(img, augment=opt.augment)[0]
 
         # Apply NMS
         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)

diff --git a/models/yolo.py b/models/yolo.py
@@ -59,7 +59,7 @@ def forward(self, x):
                     y = torch.cat((xy, wh, y[..., 4:]), -1)
                 z.append(y.view(bs, -1, self.no))
 
-        return x if self.training else torch.cat(z, 1)
+        return x if self.training else (torch.cat(z, 1), x)
 
     @staticmethod
     def _make_grid(nx=20, ny=20):