def METHOD_NAME(self, dataset, loader_kwargs):
    sampler = loader_kwargs.get("sampler", None)
    # Shuffling should really only matter for the train stage. Shuffling
    # will also lead to more padding in batches if the order was otherwise
    # sorted by length.
    shuffle = loader_kwargs.get("shuffle", False)
    if shuffle and not self.distributed_launch:
        if sampler is not None:
            raise ValueError("Cannot specify both shuffle=True"
                             "and a sampler in loader_kwargs")
        sampler = ReproducibleRandomSampler(dataset)
        self.train_sampler = sampler
        loader_kwargs["sampler"] = self.train_sampler
        # Delete the shuffle flag, since you cannot specify both a sampler and
        # shuffling:
        del loader_kwargs["shuffle"]
    # Possibly make a DistributedSampler or a wrapper for some other sampler
    if self.distributed_launch and not isinstance(dataset, IterableDataset):
        drop_last = loader_kwargs.get("drop_last", False)
        # num_replicas arg is equal to world_size
        # and retrieved automatically within
        # DistributedSampler obj.
        if sampler is not None:
            self.train_sampler = DistributedSamplerWrapper(
                sampler,
                rank=self.rank,
                drop_last=drop_last,
                shuffle=shuffle, )
            # with DistributedSamplerWrapper, one must disable shuffling for dataloader
            loader_kwargs["shuffle"] = False
            loader_kwargs["sampler"] = self.train_sampler
        elif loader_kwargs.get("batch_sampler") is None:
            # no sampler and batch-sampler
            self.train_sampler = DistributedSampler(
                dataset, rank=self.rank, shuffle=True, drop_last=drop_last)
            # with DistributedSamplerWrapper, one must disable shuffling for dataloader
            loader_kwargs["shuffle"] = False
            loader_kwargs["sampler"] = self.train_sampler
        else:  # batch_sampler was specified
            self.train_sampler = DistributedSamplerWrapper(
                loader_kwargs.get("batch_sampler", None),
                rank=self.rank,
                shuffle=True, )
            loader_kwargs["batch_sampler"] = self.train_sampler
    elif self.distributed_launch and isinstance(dataset, IterableDataset):
        logger.warning("Cannot automatically solve distributed sampling "
                       "for IterableDataset.")
    return loader_kwargs