From 20bb5f0a6a72eead188b91021bab0c58df70671c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 24 Apr 2019 16:39:56 +0200 Subject: [PATCH] updates --- utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/utils.py b/utils/utils.py index d5917dc6..2e46fbbe 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -351,7 +351,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5): pred[:, 4] *= class_conf # Select only suitable predictions - i = (pred[:, 4] > conf_thres) & (pred[:, 2:4] > min_wh).all(1) & (torch.isnan(pred).any(1) == 0) + i = (pred[:, 4] > conf_thres) & (pred[:, 2:4] > min_wh).all(1) & torch.isfinite(pred).all(1) pred = pred[i] # If none are remaining => process next image