Hi
I have some large text stored in TFDS format, I want to run seq2seq models on them efficiently over TPUs/multiple GPUs, in the datasets with random access, one can use the following function:
def get_tpu_sampler(dataset: torch.utils.data.dataset.Dataset):
    if xm.xrt_world_size() <= 1:
        return RandomSampler(dataset)
    return DistributedSampler(dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())

I am not sure how to write dataloaders allowing distributed training over TPUs when datasets are iterable. Could you provide me with examples and best way to handle large-scale datasets in pytorch?
thanks