polydefkis commited on
Commit
1fc9d42
·
unverified ·
1 Parent(s): 8bc0027

NMS --classes 0 bug fix (#1710)

Browse files
Files changed (1) hide show
  1. utils/general.py +1 -1
utils/general.py CHANGED
@@ -311,7 +311,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
311
  x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
312
 
313
  # Filter by class
314
- if classes:
315
  x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
316
 
317
  # Apply finite constraint
 
311
  x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
312
 
313
  # Filter by class
314
+ if classes is not None:
315
  x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
316
 
317
  # Apply finite constraint