examples/flax/language-modeling/run_mlm_flax.py/write_eval_metric
def write_eval_metric(summary_writer, eval_metrics, step):
    for metric_name, value in eval_metrics.items():
        summary_writer.scalar(f"eval_{metric_name}", value, step)