class_docstr: Vision Transformer as per https://arxiv.org/abs/2010.11929.
class_docstr_tok: ["Vision","Transformer","as","per","https",":","\/\/arxiv.org\/abs\/2010.11929","."]
code: def forward(self, x: torch.Tensor):
      assert x.ndim == 4, "Unexpected input shape"
      n, c, h, w = x.shape
      p = self.patch_size
      assert h == w == self.image_size
      n_h = h // p
      n_w = w // p
  
      # (n, c, h, w) -> (n, hidden_dim, n_h, n_w)
      x = self.conv_proj(x)
      # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w))
      x = x.reshape(n, self.hidden_dim, n_h * n_w)
  
      # (n, hidden_dim, (n_h * n_w)) -> ((n_h * n_w), n, hidden_dim)
      # the self attention layer expects inputs in the format (S, N, E)
      # where S is the source sequence length, N is the batch size, E is the
      # embedding dimension
      x = x.permute(2, 0, 1)
  
      if self.classifier == "token":
          # expand the class token to the full batch
          batch_class_token = self.class_token.expand(-1, n, -1)
          x = torch.cat([batch_class_token, x], dim=0)
  
      x = self.encoder(x)
  
      if self.classifier == "token":
          # just return the output for the class token
          x = x[0, :, :]
      else:
          x = x.mean(dim=0)
  
      return self.trunk_output(x)
  
imports: import copy
  import logging
  import math
  from collections import OrderedDict
  from functools import partial
  from typing import Dict, List, Mapping, NamedTuple, Union
  import torch
  import torch.nn
  from classy_vision.models import ClassyModel, register_model
  from .lecun_normal_init import lecun_normal_init
label_desc: Use image node_id 4 for calling the VisionTransformer obj's underlying member method code with example usage: obj.forward(x) and returns: self