YuriRibeiro
commited on
Commit
•
7d78b17
1
Parent(s):
a209a32
Fix utils.py nonzero() warning (#551)
Browse files- utils/utils.py +1 -1
utils/utils.py
CHANGED
@@ -599,7 +599,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
|
|
599 |
|
600 |
# Detections matrix nx6 (xyxy, conf, cls)
|
601 |
if multi_label:
|
602 |
-
i, j = (x[:, 5:] > conf_thres).nonzero().t()
|
603 |
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
604 |
else: # best class only
|
605 |
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
|
599 |
|
600 |
# Detections matrix nx6 (xyxy, conf, cls)
|
601 |
if multi_label:
|
602 |
+
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
|
603 |
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
604 |
else: # best class only
|
605 |
conf, j = x[:, 5:].max(1, keepdim=True)
|