NMS --classes 0 bug fix (#1710)

This commit is contained in:
Polydefkis Gkagkos authored and GitHub committed 2020-12-16 08:58:51 -08:00
1 parent 8bc0027afc
commit 1fc9d42a64
1 file changed
+1 -1
+1 -1
View File
@@ -311,7 +311,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
# Filter by class
if classes:
if classes is not None:
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
# Apply finite constraint