YuriRibeiro commited on
Commit
7d78b17
1 Parent(s): a209a32

Fix utils.py nonzero() warning (#551)

Browse files
Files changed (1) hide show
  1. utils/utils.py +1 -1
utils/utils.py CHANGED
@@ -599,7 +599,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
599
 
600
  # Detections matrix nx6 (xyxy, conf, cls)
601
  if multi_label:
602
- i, j = (x[:, 5:] > conf_thres).nonzero().t()
603
  x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
604
  else: # best class only
605
  conf, j = x[:, 5:].max(1, keepdim=True)
 
599
 
600
  # Detections matrix nx6 (xyxy, conf, cls)
601
  if multi_label:
602
+ i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).t()
603
  x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
604
  else: # best class only
605
  conf, j = x[:, 5:].max(1, keepdim=True)