correct the size of returned tensor in comments of ps_roi_pool.py and ps_roi_align.py (#3849)

* correct the size of returne tensor

* Styling changes

Co-authored-by: Vasilis Vryniotis <[email protected]>

Author: chengjuzhou

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)

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)