Add logging with MLFlow in train.py

[noraabiakar]

Summary

This PR adds MLflow tracking to the training pipeline, complementing the existing WandB tracking.

Changes

• Added MLflow and WandB initialization in the new init_tracking context manager
• Created a helper function mlflow_log_params_and_set_tags to set up MLflow with relevant configuration parameters
• Added log_metrics function to send metrics to both tracking systems when enabled
• Unlike WandB which returns a run object directly, MLflow uses Python's context manager pattern (with mlflow.start_run() as run). The new init_tracking function manages both approaches, yielding both run objects to the caller while ensuring MLflow's context is properly entered and exited.

Addresses #529