Ways to hyperparameter tuning runs in hydra --multirun as child runs in mlflow?

[tmdqja75]

Hello, I'm trying to perform hyperparameter tuning for deep learning experiments with hydra's multi run option.

python main.py --multirun 'model.hidden_size=choice(32, 64, 128)'

I'm looking for ways to log each hyperparmeter tuning runs as mlflow child runs.

I've tried the following and each hyperparameter tuning runs are logged as separate parent-child run pair.

@hydra.main(config_path="config", config_name="config")
def train(cfg: DictConfig):
    ....
    mlflow.set_experiment(experiment_name=cfg.log.experiment_name_index+'_'+zone_id)
    mlflow.pytorch.autolog()

    with mlflow.start_run(run_name="Parent"):
        with mlflow.start_run(run_name="Child", nested=True) as nested_run:
           model.fit(
                ...
                )
            metric = rmse(test_target, test_pred)
    ....

    return metric

here is the multirun setup in config.yaml

defaults:
  - override hydra/sweeper: optuna
...
hydra:
  output_subdir: null  
  run:  
    dir: .
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    storage: null
    n_trials: 5
    n_jobs: 1

Is there ways to log tuning runs as child runs in mlflow?