code: def image_rmse(
      predictions, labels, weights_fn=common_layers.weights_all
  ):
      """RMSE but will argmax if last dim is not 1."""
      if common_layers.shape_list(predictions)[-1] == 1:
          predictions = tf.squeeze(predictions, axis=[-1])
      else:
          predictions = tf.argmax(predictions, axis=-1)
      return padded_rmse(predictions, labels, weights_fn)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import six
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import bleu_hook
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import rouge
  from tensor2tensor.utils import sari_hook
  import tensorflow.compat.v1
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 1 for calling a global function with example usage: image_rmse(predictions, labels, weights_fn) and returns: padded_rmse