[fbsync] Update vision_transformer.py (#5820)

Summary:
the assert msg should be same

Reviewed By: jdsgomes, NicolasHug

Differential Revision: D36095680

fbshipit-source-id: 590b6befa239e06b077778886c058da1c13be550

Co-authored-by: Vasilis Vryniotis <[email protected]>

Author: 2 people

torchvision/models/vision_transformer.py

@@ -79,7 +79,7 @@ def __init__(
         self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
 
     def forward(self, input: torch.Tensor):
-        torch._assert(input.dim() == 3, f"Expected (seq_length, batch_size, hidden_dim) got {input.shape}")
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
         x = self.ln_1(input)
         x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
         x = self.dropout(x)