Update modeling_vit.py

modeling_vit.py

@@ -924,10 +924,11 @@ class ViTForSemanticSegmentation(ViTPreTrainedModel):
         out = self.decoder_norm(out)
         B, _, C = out.shape
         # out = out.reshape(B, self.hw_shape, self.hw_shape, C).permute(0, 3, 1, 2).contiguous()
-        out = out.reshape([B, self.hw_shape, self.hw_shape, C])       
+        out = out.reshape(B, self.hw_shape, self.hw_shape, C)
         out = self.decoder_mlp(out)
+        out = out.permute(0, 3, 1, 2).contiguous()
         out = nn.functional.interpolate(out, scale_factor=4, mode='bilinear', align_corners=False)
-        
+        out = out.permute(0, 2, 3, 1).contiguous()
         logits = self.decoder_classifier(out)
         logits = logits.permute(0, 3, 1, 2).contiguous()