code: def __init__(self, layers=12, embed_dim=192, num_heads=3):
      self.embedding = (
          Tensor.uniform(embed_dim, 3, 16, 16),
          Tensor.zeros(embed_dim),
      )
      self.embed_dim = embed_dim
      self.cls = Tensor.ones(1, 1, embed_dim)
      self.pos_embedding = Tensor.ones(1, 197, embed_dim)
      self.tbs = [
          TransformerBlock(
              embed_dim=embed_dim,
              num_heads=num_heads,
              ff_dim=embed_dim * 4,
              prenorm=True,
              act=lambda x: x.gelu(),
          )
          for i in range(layers)
      ]
      self.encoder_norm = (
          Tensor.uniform(embed_dim),
          Tensor.zeros(embed_dim),
      )
      self.head = (Tensor.uniform(embed_dim, 1000), Tensor.zeros(1000))
  
file_path: tinygrad/extra/models/vit.py
imports: import numpy
  from tinygrad.tensor import Tensor
  from tinygrad.helpers import fetch
  from extra.models.transformer import TransformerBlock
label_desc: Use image node_id 1 to create a new ViT object  with example: obj = ViT(layers, embed_dim, num_heads)