def METHOD_NAME(dataset, current_gpu_index, num_gpu, batch_size):
    sampler = DistributedSampler(
        dataset,
        num_replicas=num_gpu,
        rank=current_gpu_index,
        shuffle=False,
    )
    # Create a DataLoader
    train_loader = DataLoader(
        dataset,
        sampler=sampler,
        batch_size=batch_size,
        shuffle=False,
    )
    return train_loader