code: def tune_mnist(
      num_samples=10,
      num_epochs=10,
      gpus_per_trial=0,
      tracking_uri=None,
      experiment_name="ptl_autologging_example",
  ):
      data_dir = os.path.join(tempfile.gettempdir(), "mnist_data_")
      # Download data
      MNISTDataModule(data_dir=data_dir, batch_size=32).prepare_data()
  
      # Set the MLflow experiment, or create it if it does not exist.
      mlflow.set_tracking_uri(tracking_uri)
      mlflow.set_experiment(experiment_name)
  
      config = {
          "layer_1": tune.choice([32, 64, 128]),
          "layer_2": tune.choice([64, 128, 256]),
          "lr": tune.loguniform(1e-4, 1e-1),
          "batch_size": tune.choice([32, 64, 128]),
          "experiment_name": experiment_name,
          "tracking_uri": mlflow.get_tracking_uri(),
          "data_dir": os.path.join(
              tempfile.gettempdir(), "mnist_data_"
          ),
          "num_epochs": num_epochs,
      }
  
      trainable = tune.with_parameters(
          train_mnist_tune,
          data_dir=data_dir,
          num_epochs=num_epochs,
          num_gpus=gpus_per_trial,
      )
  
      tuner = tune.Tuner(
          tune.with_resources(
              trainable, resources={"cpu": 1, "gpu": gpus_per_trial}
          ),
          tune_config=tune.TuneConfig(
              metric="loss",
              mode="min",
              num_samples=num_samples,
          ),
          run_config=train.RunConfig(
              name="tune_mnist",
          ),
          param_space=config,
      )
      results = tuner.fit()
  
      print(
          "Best hyperparameters found were: ",
          results.get_best_result().config,
      )
  
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 2 for calling a global function with example usage: tune_mnist(num_samples, num_epochs, gpus_per_trial, tracking_uri, experiment_name) without return types