From 66b32bb4dd9e4095b78ba8578c246027de6295bb Mon Sep 17 00:00:00 2001 From: DieterTimmers <46112022+DieterTimmers@users.noreply.github.com> Date: Fri, 2 Feb 2024 06:31:27 +0800 Subject: [PATCH] Segment TensorRT `batch=3` fix (#7952) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/models/yolo/segment/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/models/yolo/segment/predict.py b/ultralytics/models/yolo/segment/predict.py index 8d8cd59a..9d7015ff 100644 --- a/ultralytics/models/yolo/segment/predict.py +++ b/ultralytics/models/yolo/segment/predict.py @@ -41,7 +41,7 @@ class SegmentationPredictor(DetectionPredictor): orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) results = [] - proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported + proto = preds[1][-1] if isinstance(preds[1], tuple) else preds[1] # tuple if PyTorch model or array if exported for i, pred in enumerate(p): orig_img = orig_imgs[i] img_path = self.batch[0][i]