diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py index 446489db989..8bbade08663 100644 --- a/torchvision/models/detection/retinanet.py +++ b/torchvision/models/detection/retinanet.py @@ -454,7 +454,7 @@ def postprocess_detections(self, head_outputs, anchors, image_shapes): return detections - def _anchors_per_level(self, features, HWA): + def _anchors_per_level(self, features: List[Tensor], HWA: int): # recover level sizes num_anchors_per_level = [x.size(2) * x.size(3) for x in features] HW = 0 diff --git a/torchvision/models/detection/ssd.py b/torchvision/models/detection/ssd.py index 0e70b840a50..12186447796 100644 --- a/torchvision/models/detection/ssd.py +++ b/torchvision/models/detection/ssd.py @@ -141,7 +141,7 @@ def __init__(self, backbone: nn.Module, num_classes: int, # used only on torchscript mode self._has_warned = False - def _anchors_per_level(self, features, HWA): + def _anchors_per_level(self, features: List[Tensor], HWA: int): # TODO: Discuss/refactor this workaround num_anchors_per_level = [x.size(2) * x.size(3) * anchors for x, anchors in zip(features, self.num_anchors)] HW = 0