def worker_dist()

in awsio/python/lib/io/s3/s3dataset.py [0:0]


    def worker_dist(self, urls):
        if self.dist:
            total_size = len(urls)
            urls = urls[self.rank:total_size:self.world_size]

        worker_info = torch.utils.data.get_worker_info()
        if worker_info is not None:
            wid = worker_info.id
            num_workers = worker_info.num_workers
            length = len(urls)
            return urls[wid:length:num_workers]
        else:
            return urls