examples/demo_torch_multi_gpu.py/prepare
def prepare(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