updates

2019-12-25 14:55:11 -08:00 · 2019-12-25 14:55:11 -08:00 · d1087f4987
parent 34d9392bac
commit d1087f4987
1 changed files with 2 additions and 1 deletions
--- 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':