updates
This commit is contained in:
parent
34d9392bac
commit
d1087f4987
|
@ -537,6 +537,7 @@ 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)
|
pred = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)
|
||||||
|
|
||||||
# Apply finite constraint
|
# Apply finite constraint
|
||||||
|
if not torch.isfinite(pred).all():
|
||||||
pred = pred[torch.isfinite(pred).all(1)]
|
pred = pred[torch.isfinite(pred).all(1)]
|
||||||
|
|
||||||
# Batched NMS
|
# Batched NMS
|
||||||
|
|
Loading…
Reference in New Issue