def make_image_wrapped_metric_fn(metric_fn):
    """Metric fn without tf.metrics.mean."""

    def image_wrapped_metric_fn(predictions,
                                features,
                                labels,
                                weights_fn=common_layers.weights_all):
      del weights_fn
      del features
      predictions, labels = reduce_dimensions(predictions, labels)
      return metric_fn(predictions, labels, model_hparams)

    return image_wrapped_metric_fn
Metric fn without tf.metrics.mean.
tensor2tensor/utils/metrics.py