code: def testPostProcessImageTrainMode(
      self, likelihood, num_mixtures, depth
  ):
      batch = 1
      rows = 8
      cols = 24
      hparams = hparam.HParams(
          hidden_size=2,
          likelihood=likelihood,
          mode=tf_estimator.ModeKeys.TRAIN,
          num_mixtures=num_mixtures,
      )
      inputs = tf.random_uniform(
          [batch, rows, cols, hparams.hidden_size],
          minval=-1.0,
          maxval=1.0,
      )
      outputs = common_image_attention.postprocess_image(
          inputs, rows, cols, hparams
      )
      self.assertEqual(outputs.shape, (batch, rows, cols, depth))
  
file_path: tensor2tensor/tensor2tensor/layers/common_image_attention_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.utils import hparam
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling the CommonImageAttentionTest obj's underlying member method code with example usage: obj.testPostProcessImageTrainMode(likelihood, num_mixtures, depth) without return types