Mlflow and KerasTuner integration

Question

I am trying to integrate together KerasTuner and Mlflow. I\'d like to record the loss at each epoch of each trial of Keras Tuner.

My approach is