def test_combined_data_loader_validation_test(use_distributed_sampler):
    """This test makes sure distributed sampler has been properly injected in dataloaders when using CombinedLoader."""

    class CustomDataset(Dataset):
        def __init__(self, data):
            self.data = data

        def __len__(self):
            return len(self.data)

        def __getitem__(self, index):
            return self.data[index]

    class CustomSampler(RandomSampler):
        def __init__(self, data_source, name) -> None:
            super().__init__(data_source)
            self.name = name

    dataset = CustomDataset(range(10))
    combined_loader = CombinedLoader({
        "a": DataLoader(CustomDataset(range(10))),
        "b": DataLoader(dataset, sampler=CustomSampler(dataset, "custom_sampler")),
        "c": {"c": DataLoader(CustomDataset(range(10))), "d": DataLoader(CustomDataset(range(10)))},
        "d": [DataLoader(CustomDataset(range(10))), DataLoader(CustomDataset(range(10)))],
    })
    model = BoringModel()
    trainer = Trainer(use_distributed_sampler=use_distributed_sampler, strategy="ddp", accelerator="cpu", devices=2)
    trainer.strategy.connect(model)
    trainer._data_connector.attach_data(model, train_dataloaders=combined_loader)
    trainer.fit_loop.setup_data()

    samplers_flattened = tree_flatten(combined_loader.sampler)[0]
    assert len(samplers_flattened) == 6
    if use_distributed_sampler:
        assert all(isinstance(s, DistributedSampler) for s in samplers_flattened)
    else:
        assert all(isinstance(s, (SequentialSampler, CustomSampler)) for s in samplers_flattened)

    datasets_flattened = [dl.dataset for dl in combined_loader.flattened]
    assert len(datasets_flattened) == 6
    assert all(isinstance(ds, CustomDataset) for ds in datasets_flattened)
This test makes sure distributed sampler has been properly injected in dataloaders when using CombinedLoader.
tests/tests_pytorch/utilities/test_combined_loader.py