from torch.utils.data import DataLoader, DistributedSampler


def create_loader(dataset, batch_size, rank, world_size, num_workers: int = 2, drop_last: bool = False, shuffle: bool = True):
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, drop_last=drop_last, shuffle=shuffle)

    loader_kwargs = dict(
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=(num_workers > 0),
        drop_last=drop_last,
    )

    if num_workers > 0:
        loader_kwargs["prefetch_factor"] = 2

    return DataLoader(dataset, **loader_kwargs)
