code: def load_from_pretrained(m):
      # https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
      if m.embed_dim == 192:
          url = "https://storage.googleapis.com/vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz"
      elif m.embed_dim == 768:
          url = "https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz"
      else:
          raise Exception("no pretrained weights for configuration")
      dat = np.load(fetch(url))
  
      # for x in dat.keys():
      #  print(x, dat[x].shape, dat[x].dtype)
  
      m.embedding[0].assign(
          np.transpose(dat["embedding/kernel"], (3, 2, 0, 1))
      )
      m.embedding[1].assign(dat["embedding/bias"])
  
      m.cls.assign(dat["cls"])
  
      m.head[0].assign(dat["head/kernel"])
      m.head[1].assign(dat["head/bias"])
  
      m.pos_embedding.assign(
          dat["Transformer/posembed_input/pos_embedding"]
      )
      m.encoder_norm[0].assign(dat["Transformer/encoder_norm/scale"])
      m.encoder_norm[1].assign(dat["Transformer/encoder_norm/bias"])
  
      for i in range(12):
          m.tbs[i].query[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel"
              ].reshape(m.embed_dim, m.embed_dim)
          )
          m.tbs[i].query[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/bias"
              ].reshape(m.embed_dim)
          )
          m.tbs[i].key[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/kernel"
              ].reshape(m.embed_dim, m.embed_dim)
          )
          m.tbs[i].key[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/bias"
              ].reshape(m.embed_dim)
          )
          m.tbs[i].value[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/kernel"
              ].reshape(m.embed_dim, m.embed_dim)
          )
          m.tbs[i].value[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/bias"
              ].reshape(m.embed_dim)
          )
          m.tbs[i].out[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/kernel"
              ].reshape(m.embed_dim, m.embed_dim)
          )
          m.tbs[i].out[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/bias"
              ].reshape(m.embed_dim)
          )
          m.tbs[i].ff1[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/kernel"
              ]
          )
          m.tbs[i].ff1[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/bias"
              ]
          )
          m.tbs[i].ff2[0].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/kernel"
              ]
          )
          m.tbs[i].ff2[1].assign(
              dat[
                  f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/bias"
              ]
          )
          m.tbs[i].ln1[0].assign(
              dat[f"Transformer/encoderblock_{i}/LayerNorm_0/scale"]
          )
          m.tbs[i].ln1[1].assign(
              dat[f"Transformer/encoderblock_{i}/LayerNorm_0/bias"]
          )
          m.tbs[i].ln2[0].assign(
              dat[f"Transformer/encoderblock_{i}/LayerNorm_2/scale"]
          )
          m.tbs[i].ln2[1].assign(
              dat[f"Transformer/encoderblock_{i}/LayerNorm_2/bias"]
          )
  
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 4 for calling the ViT obj's underlying member method code with example usage: obj.load_from_pretrained(m) without return types