Skip to content

Commit eb3129b

Browse files
authored
Fix _out_channels for universal timm encoder
1 parent ad3e5c1 commit eb3129b

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

segmentation_models_pytorch/encoders/timm_universal.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def __init__(self, name, pretrained=True, in_channels=3, depth=5, output_stride=
2121
self.model = timm.create_model(name, **kwargs)
2222

2323
self._in_channels = in_channels
24-
self._out_channels = [3, ] + self.model.feature_info.channels()
24+
self._out_channels = [in_channels, ] + self.model.feature_info.channels()
2525
self._depth = depth
2626

2727
def forward(self, x):

0 commit comments

Comments
 (0)