Skip to content

Commit

Permalink
minor fix
Browse files Browse the repository at this point in the history
AlexeyAB committed Jul 21, 2022

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 894a93e commit 2596994
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion detect.py
Original file line number Diff line number Diff line change
@@ -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)
2 changes: 1 addition & 1 deletion models/yolo.py
Original file line number Diff line number Diff line change
@@ -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):

0 comments on commit 2596994

Please sign in to comment.