Multitask Loop?

[micklemouse]

Hi all, I'm new to Lightning and I've been reading the docs for a couple of hours already, but I've not found what I'm looking for yet.

I'm interested in the feasibility of a multitask loop. I'm looking for something that loops like the following, while retaining the maximum amount of Lightning functionality:

for iteration in range(max_iteration):
  for (model, train_loader, val_loaders, ...) in list_of_tasks:
    # Like using yield instead of return after running TrainingEpochLoop
    trainer.fit_single_batch(model, train_loader, val_loaders, ...)

Where the models have shared weights. For example, a model backbone gets trains for one step, then we train with a head attached to the model backbone on another dataset for one step, etc. The dataset return signatures can vary wildly, sometimes requiring custom collate functions, etc.

I have read #1502 but that solution doesn't seem as to work for my case.

I have read https://stackoverflow.com/questions/71774659/combine-multiple-dataloaders-sequentially#:~:text=So%2C%20how%20would%20I%20go%20about%20combining%20multiple,something%20like%20def%20train_dataloader%20%28self%29%3A%20return%20%5Bdata_loader_1%2C%20data_loader_2%5D which is closer to what I want. I think this approach may work but will require overriding multiple loops.

What does the community recommend?