Updated annotation to be a Union of Tensor and List (#4416)

* Updated annotation to be a Union of Tensor and List

* Updated check_roi_boxes_shape.

Author: prabhat00155

torchvision/ops/_utils.py

@@ -1,6 +1,6 @@
 import torch
 from torch import Tensor
-from typing import List
+from typing import List, Union
 
 
 def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor:
@@ -24,7 +24,7 @@ def convert_boxes_to_roi_format(boxes: List[Tensor]) -> Tensor:
     return rois
 
 
-def check_roi_boxes_shape(boxes: Tensor):
+def check_roi_boxes_shape(boxes: Union[Tensor, List[Tensor]]):
     if isinstance(boxes, (list, tuple)):
         for _tensor in boxes:
             assert _tensor.size(1) == 4, \

torchvision/ops/roi_align.py

@@ -1,16 +1,17 @@
+from typing import List, Union
+
 import torch
 from torch import nn, Tensor
-
 from torch.nn.modules.utils import _pair
 from torch.jit.annotations import BroadcastingList2
-
 from torchvision.extension import _assert_has_ops
+
 from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape
 
 
 def roi_align(
     input: Tensor,
-    boxes: Tensor,
+    boxes: Union[Tensor, List[Tensor]],
     output_size: BroadcastingList2[int],
     spatial_scale: float = 1.0,
     sampling_ratio: int = -1,

torchvision/ops/roi_pool.py

@@ -1,16 +1,17 @@
+from typing import List, Union
+
 import torch
 from torch import nn, Tensor
-
 from torch.nn.modules.utils import _pair
 from torch.jit.annotations import BroadcastingList2
-
 from torchvision.extension import _assert_has_ops
+
 from ._utils import convert_boxes_to_roi_format, check_roi_boxes_shape
 
 
 def roi_pool(
     input: Tensor,
-    boxes: Tensor,
+    boxes: Union[Tensor, List[Tensor]],
     output_size: BroadcastingList2[int],
     spatial_scale: float = 1.0,
 ) -> Tensor: