class_docstr: The search space that is proposed in `AutoFormer <https://arxiv.org/abs/2107.00651>`__.
  There are four searchable variables: depth, embedding dimension, heads number and MLP ratio.
  
  Parameters
  ----------
  search_embed_dim
      The search space of embedding dimension. Use a list to specify search range.
  search_mlp_ratio
      The search space of MLP ratio. Use a list to specify search range.
  search_num_heads
      The search space of number of heads. Use a list to specify search range.
  search_depth
      The search space of depth. Use a list to specify search range.
  img_size
      Size of input image.
  patch_size
      Size of image patch.
  in_channels
      Number of channels of the input image.
  num_labels
      Number of classes for classifier.
  qkv_bias
      Whether to use bias item in the qkv embedding.
  drop_rate
      Drop rate of the MLP projection in MSA and FFN.
  attn_drop_rate
      Drop rate of attention.
  drop_path_rate
      Drop path rate.
  pre_norm
      Whether to use pre_norm. Otherwise post_norm is used.
  global_pooling
      Whether to use global pooling to generate the image representation. Otherwise the cls_token is used.
  absolute_position
      Whether to use absolute positional embeddings.
  qk_scale
      The scaler on score map in self-attention.
  rpe
      Whether to use relative position encoding.
docstr: Load the searched subnet model.
  
  Parameters
  ----------
  name
      Search space size, must be one of {'autoformer-tiny', 'autoformer-small', 'autoformer-base'}.
  pretrained
      Whether initialized with pre-trained weights.
  download
      Whether to download supernet weights.
  progress
      Whether to display the download progress.
  
  Returns
  -------
  nn.Module
      The subnet model.
docstr_tok: ["Load","the","searched","subnet","model",".","Parameters","--","--","--","--","--","name","Search","space","size",",","must","be","one","of","{","'autoformer-tiny","'",",","'autoformer-small","'",",","'autoformer-base","'","}",".","pretrained","Whether","initialized","with","pre-trained","weights",".","download","Whether","to","download","supernet","weights",".","progress","Whether","to","display","the","download","progress",".","Returns","--","--","--","-","nn.Module","The","subnet","model","."]
code: def load_searched_model(
      cls,
      name: str,
      pretrained: bool = False,
      download: bool = True,
      progress: bool = True,
  ) -> nn.Module:
      """
      Load the searched subnet model.
  
      Parameters
      ----------
      name
          Search space size, must be one of {'autoformer-tiny', 'autoformer-small', 'autoformer-base'}.
      pretrained
          Whether initialized with pre-trained weights.
      download
          Whether to download supernet weights.
      progress
          Whether to display the download progress.
  
      Returns
      -------
      nn.Module
          The subnet model.
      """
      legal = ["autoformer-tiny", "autoformer-small", "autoformer-base"]
      if name not in legal:
          raise ValueError(
              f"Unsupported name: {name}. It should be one of {legal}."
          )
      name = name[11:]
      init_kwargs = cls.preset(name)
      if name == "tiny":
          mlp_ratio = [
              3.5,
              3.5,
              3.0,
              3.5,
              3.0,
              3.0,
              4.0,
              4.0,
              3.5,
              4.0,
              3.5,
              4.0,
              3.5,
          ] + [3.0]
          num_head = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3] + [3]
          arch: Dict[str, Any] = {"embed_dim": 192, "depth": 13}
          for i in range(14):
              arch[f"mlp_ratio_{i}"] = mlp_ratio[i]
              arch[f"num_head_{i}"] = num_head[i]
      elif name == "small":
          mlp_ratio = [
              3.0,
              3.5,
              3.0,
              3.5,
              4.0,
              4.0,
              4.0,
              4.0,
              4.0,
              4.0,
              4.0,
              3.5,
              4.0,
          ] + [3.0]
          num_head = [6, 6, 5, 7, 5, 5, 5, 6, 6, 7, 7, 6, 7] + [5]
          arch: Dict[str, Any] = {"embed_dim": 384, "depth": 13}
          for i in range(14):
              arch[f"mlp_ratio_{i}"] = mlp_ratio[i]
              arch[f"num_head_{i}"] = num_head[i]
      elif name == "base":
          mlp_ratio = [
              3.5,
              3.5,
              4.0,
              3.5,
              4.0,
              3.5,
              3.5,
              3.0,
              4.0,
              4.0,
              3.0,
              4.0,
              3.0,
              3.5,
          ] + [3.0, 3.0]
          num_head = [9, 9, 9, 9, 9, 10, 9, 9, 10, 9, 10, 9, 9, 10] + [
              8,
              8,
          ]
          arch: Dict[str, Any] = {"embed_dim": 576, "depth": 14}
          for i in range(16):
              arch[f"mlp_ratio_{i}"] = mlp_ratio[i]
              arch[f"num_head_{i}"] = num_head[i]
      else:
          raise ValueError(
              f"Unsupported architecture with name: {name}"
          )
  
      model_factory = cls.frozen_factory(arch)
      model = model_factory(**init_kwargs)
  
      if pretrained:
          weight_file = load_pretrained_weight(
              f"autoformer-{name}-subnet",
              download=download,
              progress=progress,
          )
          pretrained_weights = torch.load(weight_file)
          model.load_state_dict(pretrained_weights)
  
      return model
  
imports: from __future__ import annotations
  from typing import Tuple, cast, Any, Dict
  import torch
  import torch.nn
  import torch.nn.functional
  import nni
  from nni.mutable import Categorical, MutableExpression, frozen_context, ensure_frozen, Mutable
  from nni.nas.nn.pytorch import ParametrizedModule, ModelSpace, MutableModule, Repeat, MutableConv2d, MutableLinear, MutableLayerNorm
  from nni.nas.space import current_model
  from nni.nas.oneshot.pytorch.supermodule.operation import MixedOperation
  from nni.nas.profiler.pytorch.flops import FlopsResult
  from nni.nas.profiler.pytorch.utils import MutableShape, ShapeTensor, profiler_leaf_module
  from .utils.pretrained import load_pretrained_weight
  from .utils.nn import DropPath
label_desc: Use image node_id 5 for calling the AutoFormer obj's underlying member method code with example usage: obj.load_searched_model(cls, name, pretrained, download, progress) and returns: model