import tensorflow as tf

  def make_metric_fn(metric_fn):

    def wrapped_metric_fn(logits, labels):
      num, den = metric_fn(logits, labels, weights_fn=weights_fn)
      return tf.metrics.mean(num, den)

    return wrapped_metric_fn

