Context
Currently only WandB is supported for tracking the training metrics. We are finetuning on AzureML where MLFlow is natively supported.
Proposal
I implemented the changes to support MLFlow in addition to WandB and I can PR these changes if this is something you are interested in.
You can already see the changes in this branch, specifically this commit. Let me know if you're interested in the contribution.