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)
RMSE but will argmax if last dim is not 1.
tensor2tensor/utils/metrics.py