docstr: Initialize a ViT model, takes mean of the final encoder layer outputs and returns those as features for a given image.
docstr_tok: ["Initialize","a","ViT","model",",","takes","mean","of","the","final","encoder","layer","outputs","and","returns","those","as","features","for","a","given","image","."]
code: def __init__(self) -> None:
      """
      Initialize a ViT model, takes mean of the final encoder layer outputs and returns those as features for a given image.
      """
      super().__init__()
      self.model = vit_b_16(
          weights="ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1"
      ).eval()
      self.hidden_dim = 768  # Value inferred from here: https://github.com/pytorch/vision/blob/af048198f87da11f344ffba37d6962aa78b36218/torchvision/models/vision_transformer.py#L641
      self.patch_size = 16
      self.image_size = 384  # https://github.com/pytorch/vision/blob/af048198f87da11f344ffba37d6962aa78b36218/torchvision/models/vision_transformer.py#L378
      self.class_token = nn.Parameter(
          torch.zeros(1, 1, self.hidden_dim)
      )
  
imports: from PIL.Image import Image
  from typing import Callable, NamedTuple, Optional
  import torch
  import torch.nn
  from torchvision import models
  from torchvision.transforms import transforms
  from torchvision.models import vit_b_16, EfficientNet_B4_Weights
  from torchvision.models.vision_transformer import ViT_B_16_Weights
label_desc: Use image node_id 1 to create a new ViT object from inherited base classes: torch.nn with example: obj = ViT()