diff --git a/torchvision/ops/ps_roi_align.py b/torchvision/ops/ps_roi_align.py index 8097211dfe9..d42353e2b0d 100644 --- a/torchvision/ops/ps_roi_align.py +++ b/torchvision/ops/ps_roi_align.py @@ -39,7 +39,7 @@ def ps_roi_align( ``ceil(roi_width / output_width)``, and likewise for height). Default: -1 Returns: - Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs + Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs """ _assert_has_ops() check_roi_boxes_shape(boxes) diff --git a/torchvision/ops/ps_roi_pool.py b/torchvision/ops/ps_roi_pool.py index 01cee261f77..d0331e557fd 100644 --- a/torchvision/ops/ps_roi_pool.py +++ b/torchvision/ops/ps_roi_pool.py @@ -33,7 +33,7 @@ def ps_roi_pool( the box coordinates. Default: 1.0 Returns: - Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs. + Tensor[K, C / (output_size[0] * output_size[1]), output_size[0], output_size[1]]: The pooled RoIs. """ _assert_has_ops() check_roi_boxes_shape(boxes)