Validation _process_batch error

[someuser2026]

Hi, I am getting this error after the changes made to the segment validation pipeline in prepare_Batch and process_batch. This was not happening before.

Attached are the error log and the new and old versions

Partial Error log

Logging results to /Users/manishagupta/Desktop/PhD/Code/runs/cuda/segment/imgsz_224/yolo/exp_1009_2117
Starting training for 0.15 hours...

      Epoch    GPU_mem   box_loss   seg_loss   cls_loss   dfl_loss  Instances       Size
      1/100         0G      2.815      3.344      4.062      4.086          8        224: 100% ━━━━━━━━━━━━ 2/2 0.2it/s 8.6s
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Mask(P          R      mAP50  mAP50-95): 0% ──────────── 0/1  1.5s
Traceback (most recent call last):
  File "/Users/manishagupta/miniforge3/envs/ultralytics_contrib/bin/yolo", line 7, in <module>
    sys.exit(entrypoint())
             ^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/cfg/__init__.py", line 994, in entrypoint
    getattr(model, mode)(**overrides)  # default args from model
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/model.py", line 800, in train
    self.trainer.train()
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/trainer.py", line 244, in train
    self._do_train()
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/trainer.py", line 490, in _do_train
    self.metrics, self.comp_metric = self.validate()
                                     ^^^^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/trainer.py", line 706, in validate
    metrics = self.validator(self)
              ^^^^^^^^^^^^^^^^^^^^
  File "/Users/manishagupta/miniforge3/envs/ultralytics_contrib/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/engine/validator.py", line 238, in __call__
    self.update_metrics(preds, batch)
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/models/yolo/detect/val.py", line 191, in update_metrics
    **self._process_batch(predn, pbatch),
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/models/yolo/segment/val.py", line 176, in _process_batch
    iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/manishagupta/Desktop/PhD/orig_paper_code/ultralytics/ultralytics/utils/metrics.py", line 164, in mask_iou
    intersection = torch.matmul(mask1, mask2.T).clamp_(0)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mat1 and mat2 shapes cannot be multiplied (4x3136 and 12544x29)

New _prepare_batch and _process_batch

def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
        """
        Prepare a batch for training or inference by processing images and targets.

        Args:
            si (int): Batch index.
            batch (dict[str, Any]): Batch data containing images and annotations.

        Returns:
            (dict[str, Any]): Prepared batch with processed annotations.
        """
        prepared_batch = super()._prepare_batch(si, batch)
        nl = prepared_batch["cls"].shape[0]
        if self.args.overlap_mask:
            masks = batch["masks"][si]
            index = torch.arange(1, nl + 1, device=masks.device).view(nl, 1, 1)
            masks = (masks == index).float()
        else:
            masks = batch["masks"][batch["batch_idx"] == si]
        if nl:
            mask_size = [s if self.process is ops.process_mask_native else s // 4 for s in prepared_batch["imgsz"]]
            if masks.shape[1:] != mask_size:
                masks = F.interpolate(masks[None], mask_size, mode="bilinear", align_corners=False)[0]
                masks = masks.gt_(0.5)
        prepared_batch["masks"] = masks
        return prepared_batch

    def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
        """
        Compute correct prediction matrix for a batch based on bounding boxes and optional masks.

        Args:
            preds (dict[str, torch.Tensor]): Dictionary containing predictions with keys like 'cls' and 'masks'.
            batch (dict[str, Any]): Dictionary containing batch data with keys like 'cls' and 'masks'.

        Returns:
            (dict[str, np.ndarray]): A dictionary containing correct prediction matrices including 'tp_m' for mask IoU.

        Notes:
            - If `masks` is True, the function computes IoU between predicted and ground truth masks.
            - If `overlap` is True and `masks` is True, overlapping masks are taken into account when computing IoU.

        Examples:
            >>> preds = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
            >>> batch = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
            >>> correct_preds = validator._process_batch(preds, batch)
        """
        tp = super()._process_batch(preds, batch)
        gt_cls = batch["cls"]
        if gt_cls.shape[0] == 0 or preds["cls"].shape[0] == 0:
            tp_m = np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)
        else:
            iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1))
            tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy()
        tp.update({"tp_m": tp_m})  # update tp with mask IoU
        return tp

Old versions

def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
        """
        Prepare a batch for training or inference by processing images and targets.

        Args:
            si (int): Batch index.
            batch (Dict[str, Any]): Batch data containing images and annotations.

        Returns:
            (Dict[str, Any]): Prepared batch with processed annotations.
        """
        prepared_batch = super()._prepare_batch(si, batch)
        midx = [si] if self.args.overlap_mask else batch["batch_idx"] == si
        prepared_batch["masks"] = batch["masks"][midx]
        return prepared_batch

    def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
        """
        Compute correct prediction matrix for a batch based on bounding boxes and optional masks.

        Args:
            preds (Dict[str, torch.Tensor]): Dictionary containing predictions with keys like 'cls' and 'masks'.
            batch (Dict[str, Any]): Dictionary containing batch data with keys like 'cls' and 'masks'.

        Returns:
            (Dict[str, np.ndarray]): A dictionary containing correct prediction matrices including 'tp_m' for mask IoU.

        Notes:
            - If `masks` is True, the function computes IoU between predicted and ground truth masks.
            - If `overlap` is True and `masks` is True, overlapping masks are taken into account when computing IoU.

        Examples:
            >>> preds = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
            >>> batch = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
            >>> correct_preds = validator._process_batch(preds, batch)
        """
        tp = super()._process_batch(preds, batch)
        gt_cls, gt_masks = batch["cls"], batch["masks"]
        if len(gt_cls) == 0 or len(preds["cls"]) == 0:
            tp_m = np.zeros((len(preds["cls"]), self.niou), dtype=bool)
        else:
            pred_masks = preds["masks"]
            if self.args.overlap_mask:
                nl = len(gt_cls)
                index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
                gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
                gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
            if gt_masks.shape[1:] != pred_masks.shape[1:]:
                gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
                gt_masks = gt_masks.gt_(0.5)
            iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
            tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy()
        tp.update({"tp_m": tp_m})  # update tp with mask IoU
        return tp

[UltralyticsAssistant]

👋 Hello @someuser2026, thanks for opening this discussion and for using Ultralytics 🚀. This is an automated response to help you get support faster—an Ultralytics engineer will follow up here soon.

We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it. Given your validation error in _process_batch, please include:

• Exact command you ran (CLI or Python) and model used (e.g., YOLO11 segment)
• ultralytics version/commit, OS, Python, PyTorch, and CUDA details
• A tiny sample (few images/labels) that reproduces the issue
• Full logs and stack trace
• Any custom code diffs or config overrides affecting validation or masks

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Join the Ultralytics community where it suits you best. For real-time chat, head to Discord 🎧. Prefer in-depth discussions? Check out Discourse. Or dive into threads on our Subreddit to share knowledge with the community.

Upgrade

Upgrade to the latest ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8 to verify your issue is not already resolved in the latest version:

pip install -U ultralytics

Environments

YOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[glenn-jocher]

Thanks for the detailed diffs—this is a shape mismatch: mask_iou expects both GT and prediction masks flattened from the same H×W; resize GT masks to preds["masks"].shape[1:] inside _process_batch (and drop the pre-resize in _prepare_batch) before calling IoU, e.g.:

# inside _process_batch before mask_iou()
pred_masks = preds["masks"].float()            # (N, H, W)
gt_masks   = batch["masks"].float()            # (M, h, w)

if gt_masks.shape[1:] != pred_masks.shape[1:]:
    gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0].gt_(0.5)

iou = mask_iou(gt_masks.flatten(1), pred_masks.flatten(1))
tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy()

For context on expected shapes, see the mask_iou reference in our docs: ultralytics.utils.metrics.mask_iou; if the issue persists, please paste the full text traceback and the output of:

yolo checks

[someuser2026]

Yes, that is true. This was similar to what the previous versions were, which worked perfectly fine. was there a reason to change to the current version?