code: def __init__(
      self,
      *,
      num_classes,
      s1_emb_dim=64,
      s1_emb_kernel=7,
      s1_emb_stride=4,
      s1_proj_kernel=3,
      s1_kv_proj_stride=2,
      s1_heads=1,
      s1_depth=1,
      s1_mlp_mult=4,
      s2_emb_dim=192,
      s2_emb_kernel=3,
      s2_emb_stride=2,
      s2_proj_kernel=3,
      s2_kv_proj_stride=2,
      s2_heads=3,
      s2_depth=2,
      s2_mlp_mult=4,
      s3_emb_dim=384,
      s3_emb_kernel=3,
      s3_emb_stride=2,
      s3_proj_kernel=3,
      s3_kv_proj_stride=2,
      s3_heads=6,
      s3_depth=10,
      s3_mlp_mult=4,
      dropout=0.0,
      channels=3,
  ):
      super().__init__()
      kwargs = dict(locals())
  
      dim = channels
      layers = []
  
      for prefix in ("s1", "s2", "s3"):
          config, kwargs = group_by_key_prefix_and_remove_prefix(
              f"{prefix}_", kwargs
          )
  
          layers.append(
              nn.Sequential(
                  nn.Conv2d(
                      dim,
                      config["emb_dim"],
                      kernel_size=config["emb_kernel"],
                      padding=(config["emb_kernel"] // 2),
                      stride=config["emb_stride"],
                  ),
                  LayerNorm(config["emb_dim"]),
                  Transformer(
                      dim=config["emb_dim"],
                      proj_kernel=config["proj_kernel"],
                      kv_proj_stride=config["kv_proj_stride"],
                      depth=config["depth"],
                      heads=config["heads"],
                      mlp_mult=config["mlp_mult"],
                      dropout=dropout,
                  ),
              )
          )
  
          dim = config["emb_dim"]
  
      self.layers = nn.Sequential(*layers)
  
      self.to_logits = nn.Sequential(
          nn.AdaptiveAvgPool2d(1),
          Rearrange("... () () -> ..."),
          nn.Linear(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 CvT object from inherited base classes: nn with example: obj = CvT()