code: def __init__(
      self,
      *,
      image_size,
      num_classes,
      sm_dim,
      lg_dim,
      sm_patch_size=12,
      sm_enc_depth=1,
      sm_enc_heads=8,
      sm_enc_mlp_dim=2048,
      sm_enc_dim_head=64,
      lg_patch_size=16,
      lg_enc_depth=4,
      lg_enc_heads=8,
      lg_enc_mlp_dim=2048,
      lg_enc_dim_head=64,
      cross_attn_depth=2,
      cross_attn_heads=8,
      cross_attn_dim_head=64,
      depth=3,
      dropout=0.1,
      emb_dropout=0.1,
  ):
      super().__init__()
      self.sm_image_embedder = ImageEmbedder(
          dim=sm_dim,
          image_size=image_size,
          patch_size=sm_patch_size,
          dropout=emb_dropout,
      )
      self.lg_image_embedder = ImageEmbedder(
          dim=lg_dim,
          image_size=image_size,
          patch_size=lg_patch_size,
          dropout=emb_dropout,
      )
  
      self.multi_scale_encoder = MultiScaleEncoder(
          depth=depth,
          sm_dim=sm_dim,
          lg_dim=lg_dim,
          cross_attn_heads=cross_attn_heads,
          cross_attn_dim_head=cross_attn_dim_head,
          cross_attn_depth=cross_attn_depth,
          sm_enc_params=dict(
              depth=sm_enc_depth,
              heads=sm_enc_heads,
              mlp_dim=sm_enc_mlp_dim,
              dim_head=sm_enc_dim_head,
          ),
          lg_enc_params=dict(
              depth=lg_enc_depth,
              heads=lg_enc_heads,
              mlp_dim=lg_enc_mlp_dim,
              dim_head=lg_enc_dim_head,
          ),
          dropout=dropout,
      )
  
      self.sm_mlp_head = nn.Sequential(
          nn.LayerNorm(sm_dim), nn.Linear(sm_dim, num_classes)
      )
      self.lg_mlp_head = nn.Sequential(
          nn.LayerNorm(lg_dim), nn.Linear(lg_dim, num_classes)
      )
  
imports: import torch
  from torch import nn, einsum
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 1 to create a new CrossViT object from inherited base classes: nn with example: obj = CrossViT()