code: def train_mnist_tune(
      config, data_dir=None, num_epochs=10, num_gpus=0
  ):
      setup_mlflow(
          config,
          experiment_name=config.get("experiment_name", None),
          tracking_uri=config.get("tracking_uri", None),
      )
  
      model = LightningMNISTClassifier(config, data_dir)
      dm = MNISTDataModule(
          data_dir=data_dir,
          num_workers=1,
          batch_size=config["batch_size"],
      )
      metrics = {"loss": "ptl/val_loss", "acc": "ptl/val_accuracy"}
      mlflow.pytorch.autolog()
      trainer = pl.Trainer(
          max_epochs=num_epochs,
          gpus=num_gpus,
          progress_bar_refresh_rate=0,
          callbacks=[TuneReportCallback(metrics, on="validation_end")],
      )
      trainer.fit(model, dm)
  
file_path: ray/python/ray/tune/examples/mlflow_ptl.py
imports: import os
  import tempfile
  import pytorch_lightning
  import mlflow
  from ray import train, tune
  from ray.air.integrations.mlflow import setup_mlflow
  from ray.tune.integration.pytorch_lightning import TuneReportCallback
  from ray.tune.examples.mnist_ptl_mini import LightningMNISTClassifier, MNISTDataModule
label_desc: Use image node_id 1 for calling a global function with example usage: train_mnist_tune(config, data_dir, num_epochs, num_gpus) without return types