code: def __init__(
      self,
      *,
      image_size,
      patch_size,
      num_classes,
      dim,
      depth,
      heads,
      mlp_dim,
      channels=3,
      dim_head=64,
      dropout=0.0,
      emb_dropout=0.0,
      token_dropout_prob=None,
  ):
      super().__init__()
      image_height, image_width = pair(image_size)
  
      # what percent of tokens to dropout
      # if int or float given, then assume constant dropout prob
      # otherwise accept a callback that in turn calculates dropout prob from height and width
  
      self.calc_token_dropout = None
  
      if callable(token_dropout_prob):
          self.calc_token_dropout = token_dropout_prob
  
      elif isinstance(token_dropout_prob, (float, int)):
          assert 0.0 < token_dropout_prob < 1.0
          token_dropout_prob = float(token_dropout_prob)
          self.calc_token_dropout = (
              lambda height, width: token_dropout_prob
          )
  
      # calculate patching related stuff
  
      assert divisible_by(image_height, patch_size) and divisible_by(
          image_width, patch_size
      ), "Image dimensions must be divisible by the patch size."
  
      patch_height_dim, patch_width_dim = (
          image_height // patch_size
      ), (image_width // patch_size)
      patch_dim = channels * (patch_size**2)
  
      self.channels = channels
      self.patch_size = patch_size
  
      self.to_patch_embedding = nn.Sequential(
          LayerNorm(patch_dim),
          nn.Linear(patch_dim, dim),
          LayerNorm(dim),
      )
  
      self.pos_embed_height = nn.Parameter(
          torch.randn(patch_height_dim, dim)
      )
      self.pos_embed_width = nn.Parameter(
          torch.randn(patch_width_dim, dim)
      )
  
      self.dropout = nn.Dropout(emb_dropout)
  
      self.transformer = Transformer(
          dim, depth, heads, dim_head, mlp_dim, dropout
      )
  
      # final attention pooling queries
  
      self.attn_pool_queries = nn.Parameter(torch.randn(dim))
      self.attn_pool = Attention(
          dim=dim, dim_head=dim_head, heads=heads
      )
  
      # output to logits
  
      self.to_latent = nn.Identity()
  
      self.mlp_head = nn.Sequential(
          LayerNorm(dim), nn.Linear(dim, num_classes, bias=False)
      )
  
imports: from functools import partial
  from typing import List, Union
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from torch.nn.utils.rnn import pad_sequence
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 1 to create a new NaViT object from inherited base classes: nn with example: obj = NaViT()