How to get a homogenous batch when loading data from several datasets

[talbaumel]

Hi :)

I'm trying to achieve something very similar to what is proposed in this example (from this page:

class LitModel(LightningModule):
    def train_dataloader(self):

        loader_a = torch.utils.data.DataLoader(range(8), batch_size=4)
        loader_b = torch.utils.data.DataLoader(range(16), batch_size=2)

        return {"a": loader_a, "b": loader_b}

    def training_step(self, batch, batch_idx):
        # access a dictionnary with a batch from each DataLoader
        batch_a = batch["a"]
        batch_b = batch["b"]

But instead of sampling from both datasets in each batch to randomly select a dataset and get a batch, the name of the sampled dataset must be available (since I use several unbalanced dataset with different schemas).

Thanks!