class_docstr: This class implements the `MobileViT architecture <https://arxiv.org/abs/2110.02178?context=cs.LG>`_
class_docstr_tok: ["This","class","implements","the","`","MobileViT","architecture","<","https",":","\/\/arxiv.org\/abs\/2110.02178","?","context=cs.LG",">","`","_"]
code: def _make_mit_layer(
      input_channel: int, cfg: Dict
  ) -> [nn.Sequential, int]:
      stride = cfg.get("stride", 1)
      block = []
  
      if stride == 2:
          layer = InvertedResidual(
              in_channels=input_channel,
              out_channels=cfg.get("out_channels"),
              stride=stride,
              expand_ratio=cfg.get("mv_expand_ratio", 4),
          )
  
          block.append(layer)
          input_channel = cfg.get("out_channels")
  
      transformer_dim = cfg["transformer_channels"]
      ffn_dim = cfg.get("ffn_dim")
      num_heads = cfg.get("num_heads", 4)
      head_dim = transformer_dim // num_heads
  
      if transformer_dim % head_dim != 0:
          raise ValueError(
              "Transformer input dimension should be divisible by head dimension. "
              "Got {} and {}.".format(transformer_dim, head_dim)
          )
  
      block.append(
          MobileViTBlock(
              in_channels=input_channel,
              transformer_dim=transformer_dim,
              ffn_dim=ffn_dim,
              n_transformer_blocks=cfg.get("transformer_blocks", 1),
              patch_h=cfg.get("patch_h", 2),
              patch_w=cfg.get("patch_w", 2),
              dropout=cfg.get("dropout", 0.1),
              ffn_dropout=cfg.get("ffn_dropout", 0.0),
              attn_dropout=cfg.get("attn_dropout", 0.1),
              head_dim=head_dim,
              conv_ksize=3,
          )
      )
  
      return nn.Sequential(*block), input_channel
  
imports: from typing import Optional, Tuple, Union, Dict
  import math
  import torch
  import torch.nn
  from torch import Tensor
  from torch.nn import functional
  from transformer import TransformerEncoder
  from model_config import get_config
label_desc: Use image node_id 4 for calling the MobileViT obj's underlying member method code with example usage: obj._make_mit_layer(input_channel, cfg) and returns: nn,  input_channel