def make_loader(
    dataset, sample_size=None, batch_size=1, sampler=None, random_sample=None, **kwargs
):
    """Utility for creating a dataloader on fixed sample subset."""
    import typing

    if isinstance(dataset, typing.Callable):
        # To support deferred dataset loading, support passing a factory
        # that creates the dataset when called.
        dataset = dataset()
    if isinstance(dataset, torch.Tensor):
        # The dataset can be a simple tensor.
        dataset = torch.utils.data.TensorDataset(dataset)
    if sample_size is not None:
        assert sampler is None, "sampler cannot be specified with sample_size"
        if sample_size > len(dataset):
            print(
                "Warning: sample size %d > dataset size %d"
                % (sample_size, len(dataset))
            )
            sample_size = len(dataset)
        if random_sample is None:
            sampler = FixedSubsetSampler(list(range(sample_size)))
        else:
            sampler = FixedRandomSubsetSampler(
                dataset, seed=random_sample, end=sample_size
            )
    return torch.utils.data.DataLoader(
        dataset, sampler=sampler, batch_size=batch_size, **kwargs
    )
Utility for creating a dataloader on fixed sample subset.
util/runningstats.py