Below is the model class definition

import pytorch_lightning as pl
import torch
import numpy as np
from torch.nn import MSELoss, L1Loss
from torchmetrics import R2Score

torch.random.manual_seed(123)


class LightningModelSimple(pl.LightningModule):
    def __init__(
        self,
        latent_model,
        readout_model=None,
        losses={},
        metrics=[],
        gpu=True,
        learning_rate=0.001,
        weight_decay=0.0,
    ):
        super().__init__()
        self.save_hyperparameters()
        self.latent_model = latent_model
        if readout_model is None:
            self.readout_model = torch.nn.Identity()
        else:
            self.readout_model = readout_model

        # losses
        if "target" in losses:
            self.loss_target = losses["target"]
        else:
            self.loss_target = None
        if "latent_target" in losses:
            self.loss_latent_target = losses["latent_target"]
            self.weight_loss_latent_target = losses["weight_loss_latent_target"]
        else:
            self.loss_latent_target = None

        self.gpu = gpu
        self.metrics = metrics
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay

    def forward(self, x):
        x_latent = self.latent_model(x)
        y = self.readout_model(x_latent)
        return y

    def step(self, partition, batch, batch_idx):
        spectra, target_glucose = batch

        # get latent predictions
        self.pred_latent = self.latent_model(spectra.float())

        # get glucose predictions
        self.pred_glucose = self.readout_model(self.pred_latent)

        # compute losses
        loss = 0
        if self.loss_target is not None:
            loss += self.loss_target(self.pred_glucose, target_glucose)
            self.log(partition + "_loss_target", loss, on_epoch=True)

        if self.loss_latent_target is not None:
            loss_latent_target = (
                self.weight_loss_latent_target
                * self.loss_latent_target(self.pred_latent, target_glucose.unsqueeze(1))
            )
            self.log(
                partition + "_loss_latent_target", loss_latent_target, on_epoch=True
            )
            loss += loss_latent_target

        self.log(partition + "_loss_total", loss, on_epoch=True)
        for metric_name, metric in self.metrics:
            self.log(
                partition + "_" + metric_name,
                metric(self.pred_glucose, target_glucose),
                on_epoch=True,
            )

        return loss

    def training_step(self, batch, batch_idx):
        return self.step("train", batch, batch_idx)

    def validation_step(self, batch, batch_idx):
        return self.step("val", batch, batch_idx)

    def test_step(self, batch, batch_idx):
        return self.step("test", batch, batch_idx)

    def configure_optimizers(self):
        return torch.optim.Adam(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay,
        )

This should go in a different file called helpers.py


def log_parameter(params, parser, param_name=""):
    if isinstance(params, dict):
        for key in params.keys():
            if key == "class_path":
                parser = log_parameter(params[key], parser, param_name)
            else:
                parser = log_parameter(params[key], parser, key)
    else:
        parser.add_argument("--" + param_name, type=type(params), default=params)
    return parser


def update(config_data, params):
    for k, v in params.items():
        if isinstance(v, collections.abc.Mapping):
            config_data[k] = update(config_data.get(k, {}), v)
        else:
            config_data[k] = v
    return config_data

def train_model(config_file, **kwargs):
    loader = yaml.SafeLoader
    with open(config_file, "r") as stream:
        config_data = yaml.load(stream, Loader=loader)

    if "params" in kwargs:
        config_data = update(config_data, kwargs["params"])
    if "latent_model" in kwargs:
        config_data["lightning_model"]["init_args"]["latent_model"] = kwargs[
            "latent_model"
        ]

    # experiment_name = config_data["experiment_name"]
    n_epochs = config_data["n_epochs"]

    pl.seed_everything(1234)

    # add arguments to parser
    parser = ArgumentParser(conflict_handler="resolve")
    parser.add_argument(
        "--auto-select-gpus", default=True, help="run automatically on GPU if available"
    )
    parser.add_argument("--max-epochs", default=n_epochs, type=int)
    parser.add_argument("gpus", type=int, default=1)
    parser = log_parameter(config_data, parser)

    # parse arguments to trainer
    args = parser.parse_args()

    if args.gpus == 1:
        device = "cuda"
    elif args.gpus == 0:
        device = "cpu"

    # create mlflow experiment if it doesn't yet exist
    try:
        current_experiment = dict(mlflow.get_experiment_by_name(args.experiment_name))
        experiment_id = current_experiment["experiment_id"]
    except:
        print("creating new experiment")
        experiment_id = mlflow.create_experiment(args.experiment_name)

    # # start experiment
    with mlflow.start_run(experiment_id=experiment_id) as run:
        with open("log.txt", "a") as log_file:
            log_file.write("'" + str(run.info.run_id) + "'" + ", ")

        path_mlflow_results = (
            "mlruns/" + str(experiment_id) + "/" + str(run.info.run_id)
        )
        path_checkpoints = path_mlflow_results + "/checkpoints"

        # copy yaml file to mlfow results
        # TODO: this is a hack for now, this should automatically be logged
        # with open(path_mlflow_results + "/" + config_file, "w") as f:
        with open(path_mlflow_results + "/config.yaml", "w") as f:
            yaml.dump(config_data, f)

        # initialize dataloader
        config_data = initialize_datamodule(config_data)
        datamodule = config_data["datamodule"]
        # extract key for model selection
        loss_key = config_data["metric_model_selection"]
        if (
            config_data["datamodule"].split_label_val == "Barcode"
            and "val_" in loss_key[0]
        ):
            raise ValueError(
                "split_label_val=Barcode with metric_model_selection=",
                loss_key,
                " introduces data leakage",
            )

        # initialize lightning model
        if (
            config_data["lightning_model"]["class_path"]
            == "models.lightning_model.LightningModel"
        ):
            use_val_test_data_in_train = True
        elif (
            config_data["lightning_model"]["class_path"]
            == "models.lightning_model.LightningModelSimple"
        ):
            use_val_test_data_in_train = False
        config_data = initialize_modules(config_data)
        lightning_model = config_data["lightning_model"]
        print(type(lightning_model))
        print(type(datamodule))
        # monitor different metrics depending on loss variable
        checkpoints = []
        monitored_metrics = config_data["monitored_metrics"]
        for i, (me, mo) in enumerate(monitored_metrics):
            ckpt = pl.callbacks.ModelCheckpoint(
                monitor=me,
                mode=mo,
                dirpath=path_checkpoints,
                filename="{epoch:02d}-{" + me + ":.4f}",
                save_top_k=1,
            )
            checkpoints.append(ckpt)
        # checkpoints.append(
        #     pl.callbacks.ModelCheckpoint(
        #         dirpath=path_checkpoints,
        #         filename="every_n_{epoch:02d}",
        #         every_n_epochs=10,
        #         save_top_k=-1,  # <--- this is important!
        #     )
        # )

        # log all parameter
        mlflow.pytorch.autolog()
        for arg in vars(args):
            mlflow.log_param(arg, getattr(args, arg))

        # train model
        trainer = pl.Trainer(max_epochs=n_epochs, logger=True, callbacks=checkpoints)

        # TODO: this is very hackey and should be revisited
        # we create a combined dataloader which is the same for train/validation/test
        # batching is applied to the train dataloader, thus there will be multiple batches with the batch size defined in config.yaml
        # the validation and test datloaders only have one batch which has the size of the entire validation/test set
        # insight the lightning module we read out the validation and test batch at step 0 and save it as a class
        # attribute such that all validation and test data can be used in all training steps
        if use_val_test_data_in_train:
            datamodule.setup(stage="")
            iterables_train = {
                "train": datamodule.train_dataloader(),
                "val": datamodule.val_dataloader(),
                "test": datamodule.test_dataloader(),
            }
            iterables_val = {
                "train": datamodule.train_dataloader(),
                "val": datamodule.val_dataloader(),
                "test": datamodule.test_dataloader(),
            }
            iterables_test = {
                "train": datamodule.train_dataloader(),
                "val": datamodule.val_dataloader(),
                "test": datamodule.test_dataloader(),
            }
            combined_loader_train = CombinedLoader(iterables_train, mode="max_size")
            combined_loader_val = CombinedLoader(iterables_val, mode="max_size")
            combined_loader_test = CombinedLoader(iterables_test, mode="max_size")
            trainer.fit(lightning_model, combined_loader_train, combined_loader_val)
        else:
            trainer.fit(lightning_model, datamodule=datamodule)

        # evaluate tests for all monitored metrics
        ckpts = glob.glob(path_checkpoints + "/*")
        for ckpt in ckpts:
            if loss_key[0] in ckpt:
                if use_val_test_data_in_train:
                    result = trainer.test(
                        dataloaders=combined_loader_test, ckpt_path=ckpt
                    )
                else:
                    result = trainer.test(datamodule=datamodule, ckpt_path=ckpt)
                print(result)

Finally the main file

import torch
import utils.helpers as helpers

torch.random.manual_seed(123)


if __name__ == "__main__":
    # profil data
    # train_model("config_profil_latent.yaml")
    # train_model("config_profil_readout.yaml")
    # train_model("config_profil.yaml")
    # train_model("config_profil_simple.yaml")
    for weight_decay in [1.0]:
        for val_subject in range(0, 14):
            params = {
                "datamodule": {
                    "init_args": {
                        "val_index": [val_subject],
                        "test_index": [],
                    }
                },
                "lightning_model": {
                    "init_args": {
                        "weight_decay": weight_decay,
                    }
                },
            }
            helpers.train_model("config_profil_simple.yaml", params=params)