code: def train_evaluate(params, max_epochs):
      model = IrisClassification(**params)
      dm = IrisDataModule()
      dm.setup(stage="fit")
      trainer = pl.Trainer(max_epochs=max_epochs)
      mlflow.pytorch.autolog()
      trainer.fit(model, dm)
      trainer.test(datamodule=dm)
      test_accuracy = trainer.callback_metrics.get("test_acc")
      return test_accuracy
  
imports: import argparse
  import pytorch_lightning
  from ax.service.ax_client import AxClient
  from iris import IrisClassification
  from iris_data_module import IrisDataModule
  import mlflow
label_desc: Use image node_id 1 for calling a global function with example usage: train_evaluate(params, max_epochs) and returns: test_accuracy