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