From 2596994f395a9802b0e17fbda8925c5ace793330 Mon Sep 17 00:00:00 2001 From: AlexeyAB84 Date: Thu, 21 Jul 2022 20:37:15 +0300 Subject: [PATCH] minor fix --- detect.py | 2 +- models/yolo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/detect.py b/detect.py index 841926c..53b63eb 100644 --- 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 index bdd2958..5f45aad 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -59,7 +59,7 @@ class Detect(nn.Module): 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):