From 2e1c415e5955d56c0304674bc8da80757a4c5c83 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 19 Dec 2019 20:07:58 -0800 Subject: [PATCH] updates --- utils/utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/utils/utils.py b/utils/utils.py index 6abd2b30..fc9abc67 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -501,10 +501,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru # Batched NMS if method == 'vision_batch': - i = torchvision.ops.boxes.batched_nms(boxes=pred[:, :4], - scores=pred[:, 4], - idxs=pred[:, 6], - iou_threshold=nms_thres) + i = torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres) output[image_i] = pred[i] continue