ViT Backbone Tensor Shape Fix

megatron/model/vision/vit_backbone.py

@@ -234,14 +234,20 @@ class VitBackbone(MegatronModule):
 
             token_embeddings = concatenated_tokens + \
                     self.position_embeddings(self.position_ids[:, :concatenated_tokens.shape[1]])
+            # [b, s, h] => [s, b, h]
+            token_embeddings = token_embeddings.transpose(0, 1).contiguous()
             hidden_states = self.embedding_dropout(token_embeddings)
         else:
             hidden_states = input
 
         hidden_states = self.transformer(hidden_states, None)
 
+        if self.post_process:
+            # [s b h] => [b s h]
             if self.single_token_output:
-            hidden_states = hidden_states[:,0,:]
+                hidden_states = hidden_states[0]
+            else:
+                hidden_states = hidden_states.transpose(0, 1).contiguous()
 
         return hidden_states