Use enumerate to get index of ModuleList (pytorch#4534)

Co-authored-by: Nicolas Hug <[email protected]>

Author: 2 people

torchvision/models/detection/ssd.py

@@ -62,12 +62,10 @@ def _get_result_from_module_list(self, x: Tensor, idx: int) -> Tensor:
         num_blocks = len(self.module_list)
         if idx < 0:
             idx += num_blocks
-        i = 0
         out = x
-        for module in self.module_list:
+        for i, module in enumerate(self.module_list):
             if i == idx:
                 out = module(x)
-            i += 1
         return out
 
     def forward(self, x: List[Tensor]) -> Tensor:

torchvision/ops/feature_pyramid_network.py

@@ -103,12 +103,10 @@ def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor:
         num_blocks = len(self.inner_blocks)
         if idx < 0:
             idx += num_blocks
-        i = 0
         out = x
-        for module in self.inner_blocks:
+        for i, module in enumerate(self.inner_blocks):
             if i == idx:
                 out = module(x)
-            i += 1
         return out
 
     def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor:
@@ -119,12 +117,10 @@ def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor:
         num_blocks = len(self.layer_blocks)
         if idx < 0:
             idx += num_blocks
-        i = 0
         out = x
-        for module in self.layer_blocks:
+        for i, module in enumerate(self.layer_blocks):
             if i == idx:
                 out = module(x)
-            i += 1
         return out
 
     def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]: