title: Customize metric visualization in MLFlow UI when using mlflow.tensorflow.autolog()
body: <p>I'm trying to integrate MLFlow to my project. Because I'm using <code>tf.keras.fit_generator()</code> for my training so I take advantage of <code>mlflow.tensorflow.autolog()</code>(<a href="https://www.mlflow.org/docs/latest/python_api/mlflow.tensorflow.html#mlflow.tensorflow.autolog" rel="nofollow noreferrer">docs</a> here) to enable automatic logging of metrics and parameters:</p>
  <pre><code>    model = Unet()
      optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)
  
      metrics = [IOUScore(threshold=0.5), FScore(threshold=0.5)]
      model.compile(optimizer, customized_loss, metrics)
  
      callbacks = [
          tf.keras.callbacks.ModelCheckpoint(&quot;model.h5&quot;, save_weights_only=True, save_best_only=True, mode='min'),
          tf.keras.callbacks.TensorBoard(log_dir='./logs', profile_batch=0, update_freq='batch'),
      ]
  
  
      train_dataset = Dataset(src_dir=SOURCE_DIR)
  
      train_data_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True)
  
     
      with mlflow.start_run():
          mlflow.tensorflow.autolog()
          mlflow.log_param(&quot;batch_size&quot;, BATCH_SIZE)
  
          model.fit_generator(
              train_data_loader,
              steps_per_epoch=len(train_data_loader),
              epochs=EPOCHS,
              callbacks=callbacks   
              )
  </code></pre>
  <p>I expected something like this (just a demonstration taken from the <a href="https://www.mlflow.org/docs/latest/tracking.html#visualizing-metrics" rel="nofollow noreferrer">docs</a>):</p>
  <p><a href="https://i.stack.imgur.com/eG56Z.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/eG56Z.png" alt="Visualization on the docs" /></a></p>
  <p>However, after the training finished, this is what I got:</p>
  <p><a href="https://i.stack.imgur.com/fS1JD.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/fS1JD.png" alt="f1_score visualization" /></a></p>
  <p>How can I configure so that the metric plot will update and display its value at each epoch instead of just showing the latest value?</p>