fix WongKinYiu#11

mkrupczak3 · Mar 3, 2024 · 2004346 · 2004346
1 parent 3f27231
commit 2004346
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/general.py b/utils/general.py
@@ -900,7 +900,7 @@ def non_max_suppression(
  """
 
  if isinstance(prediction, (list, tuple)): # YOLO model in validation model, output = (inference_out, loss_out)
- prediction = prediction[0] # select only inference output
+ prediction = prediction[0][1] # select only inference output
 
  device = prediction.device
  mps = 'mps' in device.type # Apple MPS