UnetSegmentationModel dimension won't match

[xdtl]

I want to try hrnet34_unet64 for image segmentation using:

encoder = E.HRNetV2Encoder34(pretrained=pretrained, layers=[0, 1, 2, 3, 4])
UnetSegmentationModel(encoder, num_classes=num_classes, unet_channels=[64, 128, 256, 512], dropout=dropout)

And got an error: ``RuntimeError: Sizes of tensors must match except in dimension 2. Got 128 and 256 (The offending index is 0)```

Could you please let me know what is wrong? Thanks!

[BloodAxe]

A HRNetV2 encoder has stem layer that outputs feature map with stride 4 (Unline resnets/densenets/effnets). UNet decoder expects that feature maps as gradually increasing dimensions. The simplest way to get it up & running - exclude this stem layer. Please take a look on the working example of the segmentation model available here: https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/zoo/segmentation.py#L92

[xdtl]

This is very helpful. Thanks so much for your reply!