Update utils/general.py

utils/general.py

@@ -990,7 +990,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][0]  # select only inference output
 
     
     device = prediction.device