From 394131c2aa66637177890a41654f75ceae61a0c8 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 20 Nov 2020 11:23:36 +0100
Subject: [PATCH] Use torchvision.ops.nms (#1460)

---
 utils/general.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils/general.py b/utils/general.py
index 2134cf1..adf766d 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -15,6 +15,7 @@ import cv2
 import matplotlib
 import numpy as np
 import torch
+import torchvision
 import yaml
 
 from utils.google_utils import gsutil_getsize
@@ -323,7 +324,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
         # Batched NMS
         c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
         boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
-        i = torch.ops.torchvision.nms(boxes, scores, iou_thres)
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
         if i.shape[0] > max_det:  # limit detections
             i = i[:max_det]
         if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)