code: def __init__(
      self,
      *,
      num_classes,
      dim,
      depth,
      heads,
      window_size=7,
      dim_head=32,
      ff_mult=4,
      channels=3,
      dropout=0.0,
  ):
      super().__init__()
      assert isinstance(
          depth, tuple
      ), "depth needs to be tuple if integers indicating number of transformer blocks at that stage"
  
      num_stages = len(depth)
  
      dims = tuple(map(lambda i: (2**i) * dim, range(num_stages)))
      dims = (channels, *dims)
      dim_pairs = tuple(zip(dims[:-1], dims[1:]))
  
      strides = (4, *((2,) * (num_stages - 1)))
  
      hyperparams_per_stage = [heads, window_size]
      hyperparams_per_stage = list(
          map(
              partial(cast_tuple, length=num_stages),
              hyperparams_per_stage,
          )
      )
      assert all(
          tuple(
              map(
                  lambda arr: len(arr) == num_stages,
                  hyperparams_per_stage,
              )
          )
      )
  
      self.layers = nn.ModuleList([])
  
      for ind, (
          (layer_dim_in, layer_dim),
          layer_depth,
          layer_stride,
          layer_heads,
          layer_window_size,
      ) in enumerate(
          zip(dim_pairs, depth, strides, *hyperparams_per_stage)
      ):
          is_last = ind == (num_stages - 1)
  
          self.layers.append(
              nn.ModuleList(
                  [
                      OverlappingPatchEmbed(
                          layer_dim_in, layer_dim, stride=layer_stride
                      ),
                      PEG(layer_dim),
                      Transformer(
                          dim=layer_dim,
                          depth=layer_depth,
                          heads=layer_heads,
                          ff_mult=ff_mult,
                          dropout=dropout,
                          norm_output=not is_last,
                      ),
                  ]
              )
          )
  
      self.mlp_head = nn.Sequential(
          Reduce("b d h w -> b d", "mean"),
          nn.LayerNorm(dims[-1]),
          nn.Linear(dims[-1], num_classes),
      )
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 1 to create a new SepViT object from inherited base classes: nn with example: obj = SepViT()