Fix type error.

pytorch · Mar 31, 2021 · eed06f4 · eed06f4
1 parent c9c8148
commit eed06f4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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
@@ -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