code: def get_model(
      hparams=None,
      mode=tf_estimator.ModeKeys.TRAIN,
      model_cls=mtf_image_transformer.MtfImageTransformer,
  ):
      if hparams is None:
          hparams = mtf_image_transformer.mtf_image_transformer_single()
      hparams.max_length = IMG_LENGTH * IMG_LENGTH
      hparams.batch_size = BATCH_SIZE
      hparams.img_len = IMG_LENGTH
      hparams.num_channels = 1
  
      p_hparams = problem_hparams.test_problem_hparams(
          VOCAB_SIZE, VOCAB_SIZE, hparams
      )
      del p_hparams.modality["inputs"]
      hparams.problem_hparams = p_hparams
  
      targets = np.random.randint(
          VOCAB_SIZE, size=(BATCH_SIZE, IMG_LENGTH, IMG_LENGTH, 1, 1)
      )
      features = {
          "targets": tf.constant(
              targets, dtype=tf.int32, name="targets"
          ),
      }
  
      return model_cls(hparams, mode, p_hparams), features, hparams
  
file_path: tensor2tensor/tensor2tensor/models/mtf_image_transformer_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import numpy
  from tensor2tensor.data_generators import problem_hparams
  from tensor2tensor.models import mtf_image_transformer
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling a global function with example usage: get_model(hparams, mode, model_cls) and returns: model_cls,  features,  hparams