Combine multiple DataLoaders sequentially
I'm interested in how I'd go about combining multiple DataLoaders sequentially for training. I understand I can use ConcatDataset to combine datasets first, but this does not work for my use case. I have a custom collate_fn that is passed to each dataloader, and this function depends on an attribute of the underlying Dataset. So, I'll have a set of custom DataLoaders like the following:
def custom_collate(sample, ref):
    data = clean_sample(torch.stack([x[0] for x in sample]), ref)
    labels = torch.tensor([x[1] for x in sample])
    return data, labels

class CollateLoader(torch.utils.data.DataLoader):
    def __init__(self, ref, *args, **kwargs):
        collate_fn = functools.partial(custom_collate, ref=ref)
        super().__init__(collate_fn = collate_fn, *args, **kwargs)

Where ref is a property of the custom Dataset class and is passed on initialization of a CollateLoader. Also, I know transforms can be applied in the Dataset, but in my case it must be done batch-wise.
So, how would I go about combining multiple DataLoaders? In the PyTorch-Lightning LightningDataModule, we can do something like
def train_dataloader(self):
    return [data_loader_1, data_loader_2]

But this will return a list of batches, not the batches sequentially.

I ran into the same problem and found a workaround. I overrided the epoch training loop using the Loops API from PytorchLightning, defining a class CustomLoop which inherits from pytorch_lightning.loops.TrainingEpochLoop, and overrided the advance() method. I copy pasted the source code from pytorch_lightning and replaced these lines with:
if not hasattr(self,'dataloader_idx'):
    self.dataloader_idx=0
if not isinstance(data_fetcher, DataLoaderIterDataFetcher):
    batch_idx = self.batch_idx + 1
    batch = next(data_fetcher.dataloader.loaders[self.dataloader_idx])
    self.dataloader_idx+=1
    if self.dataloader_idx == len(data_fetcher.dataloader.loaders):
        self.dataloader_idx = 0
else:
    batch_idx, batch = next(data_fetcher)

That way, instead of iterating over the CombinedLoader, i make it iterate over one dataloader at a time.
Then, to make use of this custom loop you have to replace the default loop in the Trainer:
trainer.fit_loop.replace(epoch_loop=CustomLoop)
trainer.fit(my_model)

