src/transformers/integrations/integration_utils.py/MLflowCallback/on_train_end
class MLflowCallback:    def on_train_end(self, args, state, control, **kwargs):
        if self._initialized and state.is_world_process_zero:
            if self._auto_end_run and self._ml_flow.active_run():
                self._ml_flow.end_run()