Updated annotation to be a Union of Tensor and List

[prabhat00155]

Resolves #3737

import torch
from torchvision.ops import roi_align, roi_pool


def _make_rois2(img_size, num_imgs, dtype, num_rois=1000): 
    rois = torch.randint(0, img_size // 2, size=(num_rois, 4)).to(dtype) 
    rois[:, 0] = torch.randint(0, num_imgs, size=(num_rois,)) 
    rois[:, 2:] += rois[:, 1:2]  
    return [rois]


pool_size = 5
img_size = 10
n_channels = 2
num_imgs = 1
dtype = torch.float

x = torch.randint(50, 100, size=(num_imgs, n_channels, img_size, img_size)).to(dtype)
rois2 = _make_rois2(img_size, num_imgs, dtype)

f1 = torch.jit.script(roi_align)
f2 = torch.jit.script(roi_pool)

print(torch.mean((f1(x, rois2, pool_size) == roi_align(x, rois2, pool_size)).to(float)))
print(torch.mean((f2(x, rois2, pool_size) == roi_pool(x, rois2, pool_size)).to(float)))

tensor(1., dtype=torch.float64)
tensor(1., dtype=torch.float64)

[datumbox]

LGTM, thanks @prabhat00155. The current op tests will fail if you break JIT scriptability so feel free to merge once the CI turns green.