Fix

qubvel-org · Jan 16, 2025 · 1ad6d80 · 1ad6d80
1 parent ebe45dc
commit 1ad6d80
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/segmentation_models_pytorch/encoders/efficientnet.py b/segmentation_models_pytorch/encoders/efficientnet.py
@@ -58,8 +58,8 @@ def __init__(
 
     def get_stages(self) -> Dict[int, Sequence[torch.nn.Module]]:
         return {
-            16: [self._blocks[self._stage_idxs[1] : self._stage_idxs[2]]],
-            32: [self._blocks[self._stage_idxs[2] :]],
+            16: [self._blocks[self._out_indexes[1] + 1 : self._out_indexes[2] + 1]],
+            32: [self._blocks[self._out_indexes[2] + 1 :]],
         }
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]: