From 48adecae5e8caea8d79b61e81d64e98ba76d9dc9 Mon Sep 17 00:00:00 2001 From: Mocha Zhao Date: Wed, 18 Jun 2025 17:05:53 +0800 Subject: [PATCH] Fix incorrect boxes slicing in postprocess_detections --- torchvision/models/detection/roi_heads.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchvision/models/detection/roi_heads.py b/torchvision/models/detection/roi_heads.py index 4e721674537..e62e34f3d54 100644 --- a/torchvision/models/detection/roi_heads.py +++ b/torchvision/models/detection/roi_heads.py @@ -698,7 +698,8 @@ def postprocess_detections( labels = labels.view(1, -1).expand_as(scores) # remove predictions with the background label - boxes = boxes[:, 1:] + boxes = boxes.reshape(-1, num_classes, 4) + boxes = boxes[..., 1:, :] scores = scores[:, 1:] labels = labels[:, 1:]