From 1fc9d42a64f248d72e1d3c7b25064998026a478f Mon Sep 17 00:00:00 2001
From: Polydefkis Gkagkos <thepolideffect@gmail.com>
Date: Wed, 16 Dec 2020 16:58:51 +0000
Subject: [PATCH] NMS --classes 0 bug fix (#1710)

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

diff --git a/utils/general.py b/utils/general.py
index 1828551..1ea7d40 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -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