code: def get_model():
      model = (
          EvoformerStack(
              c_m=256,
              c_z=128,
              c_hidden_msa_att=32,
              c_hidden_opm=32,
              c_hidden_mul=128,
              c_hidden_pair_att=32,
              c_s=384,
              no_heads_msa=8,
              no_heads_pair=4,
              no_blocks=2,  # 48
              transition_n=4,
              msa_dropout=0.15,
              pair_dropout=0.25,
              blocks_per_ckpt=None,
              inf=1000000000.0,
              eps=1e-08,
              clear_cache_between_blocks=False,
              is_multimer=False,
          )
          .eval()
          .cuda()
      )
      return model
  
file_path: colossalai/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_stack.py
imports: from typing import List, Tuple
  import pytest
  import torch
  import torch.fx
  from test_autochunk_alphafold_utils import run_test
  from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE
  from colossalai.testing import clear_cache_before_run, parameterize, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: get_model() and returns: model