Fix utils.py nonzero() warning (#551)

This commit is contained in:
YuriRibeiro authored and GitHub committed 2020-07-29 12:14:11 -07:00
1 parent a209a32019
commit 7d78b17a69
1 file changed
+1 -1
+1 -1
View File
@@ -599,7 +599,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
# Detections matrix nx6 (xyxy, conf, cls)
if multi_label:
i, j = (x[:, 5:] > conf_thres).nonzero().t()
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
else: # best class only
conf, j = x[:, 5:].max(1, keepdim=True)