code: def _test_img2img_transformer(self, net):
      batch_size = 3
      hparams = image_transformer_2d.img2img_transformer2d_tiny()
      hparams.data_dir = ""
      p_hparams = registry.problem("image_celeba").get_hparams(hparams)
      inputs = np.random.randint(256, size=(batch_size, 4, 4, 3))
      targets = np.random.randint(256, size=(batch_size, 8, 8, 3))
      with self.test_session() as session:
          features = {
              "inputs": tf.constant(inputs, dtype=tf.int32),
              "targets": tf.constant(targets, dtype=tf.int32),
              "target_space_id": tf.constant(1, dtype=tf.int32),
          }
          model = net(hparams, tf_estimator.ModeKeys.TRAIN, p_hparams)
          logits, _ = model(features)
          session.run(tf.global_variables_initializer())
          res = session.run(logits)
      self.assertEqual(res.shape, (batch_size, 8, 8, 3, 256))
  
file_path: tensor2tensor/tensor2tensor/models/image_transformer_2d_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.data_generators import celeba
  from tensor2tensor.data_generators import problem_hparams
  from tensor2tensor.models import image_transformer_2d
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling the Img2imgTransformerTest obj's underlying member method code with example usage: obj._test_img2img_transformer(net) without return types