class_docstr: Vision Transformer Model.
class_docstr_tok: ["Vision","Transformer","Model","."]
code: def forward(self, input):
      if self.pre_process:
          rearranged_input = einops.rearrange(
              input,
              "b c (h p1) (w p2) -> b (h w) (p1 p2 c)",
              p1=self.patch_dim,
              p2=self.patch_dim,
          )
  
          assert rearranged_input.dtype == torch.half
          encoder_output = self.linear_encoder(rearranged_input)
  
          concatenated_tokens = encoder_output
          if self.class_token:
              cls_tokens = self.cls_token.expand(
                  encoder_output.shape[0], -1, -1
              )
              concatenated_tokens = torch.cat(
                  (cls_tokens, encoder_output), dim=1
              )
  
          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]
          else:
              hidden_states = hidden_states.transpose(0, 1).contiguous()
  
      return hidden_states
  
file_path: megatron-lm/megatron/model/vision/vit_backbone.py
imports: import math
  import einops
  import torch
  import apex
  import torch.nn.functional
  from megatron import get_args
  from megatron.model.transformer import ParallelTransformer
  from megatron.model.utils import get_linear_layer, init_method_normal, scaled_init_method_normal
  from megatron.model.module import MegatronModule
label_desc: Use image node_id 3 for calling the VitBackbone obj's underlying member method code with example usage: obj.forward(input) and returns: hidden_states