Fixing LevelMapper.

Author: datumbox

torchvision/ops/poolers.py

@@ -40,7 +40,7 @@ def initLevelMapper(
     return LevelMapper(k_min, k_max, canonical_scale, canonical_level, eps)
 
 
-class LevelMapper(object):
+class LevelMapper(nn.Module):
     """Determine which FPN level each RoI in a set of RoIs should map to based
     on the heuristic in the FPN paper.
 
@@ -60,13 +60,14 @@ def __init__(
         canonical_level: int = 4,
         eps: float = 1e-6,
     ):
+        super().__init__()
         self.k_min = k_min
         self.k_max = k_max
         self.s0 = canonical_scale
         self.lvl0 = canonical_level
         self.eps = eps
 
-    def __call__(self, boxlists: List[Tensor]) -> Tensor:
+    def forward(self, boxlists: List[Tensor]) -> Tensor:
         """
         Args:
             boxlists (list[BoxList])
@@ -117,8 +118,7 @@ class MultiScaleRoIAlign(nn.Module):
     """
 
     __annotations__ = {
-        'scales': Optional[List[float]],
-        'map_levels': Optional[LevelMapper]
+        'scales': Optional[List[float]]
     }
 
     def __init__(
@@ -137,7 +137,7 @@ def __init__(
         self.sampling_ratio = sampling_ratio
         self.output_size = tuple(output_size)
         self.scales = None
-        self.map_levels = None
+        self.map_levels: Optional[LevelMapper] = None
         self.canonical_scale = canonical_scale
         self.canonical_level = canonical_level