def __iter__(self):
    offset = self.rank
    bins = (self.bins + self.bins[:(self.total_size - len(self.bins))])
    assert (len(bins) == self.total_size)
    samples = bins[offset::self.num_replicas]
    return iter(samples)
def __iter__(self):
    offset = self.rank
    bins = (self.bins + self.bins[:(self.total_size - len(self.bins))])
    assert (len(bins) == self.total_size)
    samples = bins[offset::self.num_replicas]
    return iter(samples)