The description of returned tensor's size in ps_roi_pool.py and ps_roi_align.py are wrong

[chengjuzhou]

📚 Documentation

The size of returned tensor should be

[K, C/(output_size[0]*output_size[1]), output_sizep[0], output_size[1]]

in
https://github.com/pytorch/vision/blob/master/torchvision/ops/ps_roi_align.py#L42
and
https://github.com/pytorch/vision/blob/master/torchvision/ops/ps_roi_pool.py#L36

This can be verified from
https://github.com/pytorch/vision/blob/master/test/test_ops.py#L186
and
https://github.com/pytorch/vision/blob/master/test/test_ops.py#L390

[datumbox]

@chengjuzhou Good spot, thanks for flagging it!

Would you like to send us a PR that fixes this?

[fmassa]

Fixed in #3849, thanks @chengjuzhou !