code: def image_rmse(predictions, labels, weights_fn=common_layers.weights_all):
    
    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)
docs: RMSE but will argmax if last dim is not 1.