fix bug for feature size of pst

zifuwan · Apr 15, 2024 · ea7c118 · ea7c118
1 parent 35abf45
commit ea7c118
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/decoders/MambaDecoder.py b/models/decoders/MambaDecoder.py
@@ -227,8 +227,11 @@ def forward_up_features(self, inputs):  # B, C, H, W
                     x = x.permute(0, 2, 3, 1).contiguous()  # B, 15, 20, 768
                     y = layer_up(x)  # B, 30, 40, 384
                 else:
+                    # interpolate y to input size (only pst900 dataset needs)
+                    B, C, H, W = inputs[3 - inx].shape
+                    y = F.interpolate(y.permute(0, 3, 1, 2).contiguous(), size=(H, W), mode='bilinear', align_corners=False).permute(0, 2, 3, 1).contiguous()
+
                     x = y + inputs[3 - inx].permute(0, 2, 3, 1).contiguous()
-
                     y = layer_up(x)
 
             x = self.norm_up(y)