Fix

qubvel · qubvel · commit 1ad6d8096fcf · 2025-01-16T16:49:52.000Z
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]:

Original file line number	Diff line number	Diff line change
`@@ -58,8 +58,8 @@ def __init__(`
`58`	`58`
`59`	`59`	`def get_stages(self) -> Dict[int, Sequence[torch.nn.Module]]:`
`60`	`60`	`return {`
`61`		`- 16: [self._blocks[self._stage_idxs[1] : self._stage_idxs[2]]],`
`62`		`- 32: [self._blocks[self._stage_idxs[2] :]],`
	`61`	`+ 16: [self._blocks[self._out_indexes[1] + 1 : self._out_indexes[2] + 1]],`
	`62`	`+ 32: [self._blocks[self._out_indexes[2] + 1 :]],`
`63`	`63`	`}`
`64`	`64`
`65`	`65`	`def forward(self, x: torch.Tensor) -> List[torch.Tensor]:`