From 2004346bcc6589a4378b2f3e662e470a4de384ad Mon Sep 17 00:00:00 2001 From: Matthew Krupczak Date: Sun, 3 Mar 2024 17:39:28 +0000 Subject: [PATCH] fix #11 --- utils/general.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/general.py b/utils/general.py index efe78b29a..af8550b3b 100644 --- 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