From d1087f49870246d3f2d56ba86c2e0b6f320273ec Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 25 Dec 2019 14:55:11 -0800
Subject: [PATCH] updates

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

diff --git a/utils/utils.py b/utils/utils.py
index 852f3d83..38ff9fd3 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -537,7 +537,8 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru
             pred = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)
 
         # Apply finite constraint
-        pred = pred[torch.isfinite(pred).all(1)]
+        if not torch.isfinite(pred).all():
+            pred = pred[torch.isfinite(pred).all(1)]
 
         # Batched NMS
         if method == 'vision_batch':