[DataLoader] Fix the world_size when distributed sharding MapDataPipe (…

…#79524)

Fixes #79449

Pull Request resolved: #79524
Approved by: https://github.com/NivekT, https://github.com/VitalyFedyunin

torch/utils/data/dataloader.py

@@ -256,6 +256,7 @@ def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1,
                 torch.utils.data.graph_settings.apply_sharding(self.dataset, ws, rank)
         elif isinstance(self.dataset, MapDataPipe):
             self.dataset = _MapDataPipeSerializationWrapper(self.dataset)
+            ws, rank = _get_distributed_settings()
             if num_workers > 0:
                 self.worker_init_fn = functools.partial(
                     _sharding_worker_init_fn, self.worker_init_fn, ws, rank)