code: def __init__(
      self,
      in_dim,
      out_dim,
      key_dim,
      downsample=("", 1),
      num_heads=8,
      attn_ratio=4,
      resolution=14,
      window_resolution=7,
      kernels=[5, 5, 5, 5],
      depth=1,
  ):
      super().__init__()
      if downsample[0] == "subsample":
          self.resolution = (resolution - 1) // downsample[1] + 1
          down_blocks = []
          down_blocks.append(
              (
                  "res1",
                  torch.nn.Sequential(
                      ResidualDrop(
                          ConvNorm(
                              in_dim, in_dim, 3, 1, 1, groups=in_dim
                          )
                      ),
                      ResidualDrop(ConvMlp(in_dim, int(in_dim * 2))),
                  ),
              )
          )
          down_blocks.append(
              ("patchmerge", PatchMerging(in_dim, out_dim))
          )
          down_blocks.append(
              (
                  "res2",
                  torch.nn.Sequential(
                      ResidualDrop(
                          ConvNorm(
                              out_dim, out_dim, 3, 1, 1, groups=out_dim
                          )
                      ),
                      ResidualDrop(ConvMlp(out_dim, int(out_dim * 2))),
                  ),
              )
          )
          self.downsample = nn.Sequential(OrderedDict(down_blocks))
      else:
          assert in_dim == out_dim
          self.downsample = nn.Identity()
          self.resolution = resolution
  
      blocks = []
      for d in range(depth):
          blocks.append(
              EfficientVitBlock(
                  out_dim,
                  key_dim,
                  num_heads,
                  attn_ratio,
                  self.resolution,
                  window_resolution,
                  kernels,
              )
          )
      self.blocks = nn.Sequential(*blocks)
  
imports: import itertools
  from collections import OrderedDict
  from typing import Dict
  import torch
  import torch.nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import SqueezeExcite, SelectAdaptivePool2d, trunc_normal_, _assert
  from ._builder import build_model_with_cfg
  from ._manipulate import checkpoint_seq
  from ._registry import register_model, generate_default_cfgs
label_desc: Use image node_id 1 to create a new EfficientVitStage object from inherited base classes: torch.nn with example: obj = EfficientVitStage(in_dim, out_dim, key_dim, downsample, num_heads, attn_ratio, resolution, window_resolution, kernels, depth)