def test_custom_batch_sampler(predicting):
    """Test that a custom (non-PyTorch) batch sampler requires the user to set `use_distributed_sampler=False`."""

    class CustomBatchSampler:  # not inheriting from `BatchSampler`
        def __iter__(self):
            while True:
                yield [0, 1, 2, 3]

    batch_sampler = CustomBatchSampler()
    dataloader = DataLoader(range(100), batch_sampler=batch_sampler)

    if predicting:
        with pytest.warns(PossibleUserWarning, match=r"Make sure your sampler is configured correctly to return all"):
            _ = _update_dataloader(dataloader, sampler=Mock(), mode=RunningStage.PREDICTING)
    else:
        with pytest.raises(TypeError, match=r"can't inject a \(distributed\) sampler into your batch sampler"):
            _ = _update_dataloader(dataloader, sampler=Mock(), mode=None)
Test that a custom (non-PyTorch) batch sampler requires the user to set `use_distributed_sampler=False`.
tests/tests_pytorch/utilities/test_data.py