Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/data_loading.py (Python)

Code before:
        is_in_dist = self.use_ddp or self.use_ddp2 or self.use_horovod or self.use_tpu
        need_dist_sampler = is_in_dist and not isinstance(dataloader.sampler, DistributedSampler)
        if self.replace_sampler_ddp and need_dist_sampler:
            if not isinstance(dataloader.sampler, (SequentialSampler, RandomSampler)):
                raise MisconfigurationException(
                    'You seem to have configured a sampler in your DataLoader. This will be replaced '
                    ' by `DistributedSampler` since `replace_sampler_ddp` is True and you are using'
                    ' distributed training. Either remove the sampler from your DataLoader or set'
                    ' `replace_sampler_ddp`=False if you want to use your custom sampler.')

            # replace with distributed sampler
            sampler = self._get_distributed_sampler(dataloader, shuffle)
            dataloader = self.replace_sampler(dataloader, sampler)

        return dataloader

    def replace_sampler(self, dataloader, sampler):
        skip_keys = ['sampler', 'batch_sampler', 'dataset_kind']

        dl_args = {
            k: v for k, v in dataloader.__dict__.items() if not k.startswith('_') and k not in skip_keys
        }

        dl_args['sampler'] = sampler
        dl_args['shuffle'] = False
        dataloader = type(dataloader)(**dl_args)
        return dataloader

    def _get_distributed_sampler(self, dataloader, shuffle):
        if self.use_tpu and XLADeviceUtils.tpu_device_exists():
            kwargs = dict(num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
        elif self.use_horovod:
            kwargs = dict(num_replicas=hvd.size(), rank=hvd.rank())
        else:
            world_size = {
                "ddp": self.num_nodes * self.num_processes,
                "ddp_spawn": self.num_nodes * self.num_processes,
                "ddp2": self.num_nodes,
                "ddp_cpu": self.num_processes * self.num_nodes
            }
            assert self.distributed_backend is not None
            kwargs = dict(num_replicas=world_size[self.distributed_backend], rank=self.global_rank)

        kwargs['shuffle'] = shuffle and not self.overfit_batches
        sampler = DistributedSampler(dataloader.dataset, **kwargs)
        return sampler

    def reset_train_dataloader(self, model: LightningModule) -> None:
        """Resets the train dataloader and initialises required variables
        (number of batches, when to validate, etc.).


Review comment (suggestion):
I think that this shall be a kind of test once at the beginning when `use_tpu` is set...

Code after:
                    ' distributed training. Either remove the sampler from your DataLoader or set'
                    ' `replace_sampler_ddp`=False if you want to use your custom sampler.')

            # replace with distributed sampler
            sampler = self._get_distributed_sampler(dataloader, shuffle)
            dataloader = self.replace_sampler(dataloader, sampler)

        return dataloader

    def replace_sampler(self, dataloader, sampler):
        skip_keys = ['sampler', 'batch_sampler', 'dataset_kind']

        dl_args = {
            k: v for k, v in dataloader.__dict__.items() if not k.startswith('_') and k not in skip_keys
        }

        dl_args['sampler'] = sampler
        dl_args['shuffle'] = False
        multiprocessing_context = dataloader.multiprocessing_context
        dataloader = type(dataloader)(**dl_args)
        dataloader.multiprocessing_context = multiprocessing_context
        return dataloader

    def _get_distributed_sampler(self, dataloader, shuffle):
        if self.use_tpu:
            kwargs = dict(num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
        elif self.use_horovod:
            kwargs = dict(num_replicas=hvd.size(), rank=hvd.rank())
        else:
            world_size = {
                "ddp": self.num_nodes * self.num_processes,
                "ddp_spawn": self.num_nodes * self.num_processes,
                "ddp2": self.num_nodes,
                "ddp_cpu": self.num_processes * self.num_nodes
            }
            assert self.distributed_backend is not None
            kwargs = dict(num_replicas=world_size[self.distributed_backend], rank=self.global_rank)

        kwargs['shuffle'] = shuffle and not self.overfit_batches
        sampler = DistributedSampler(dataloader.dataset, **kwargs)
        return sampler

    def reset_train_dataloader(self, model: LightningModule) -> None:
        """Resets the train dataloader and initialises required variables
        (number of batches, when to validate, etc.).

        Args:
            model: The current `LightningModule`
        """
        self.train_dataloader = self.request_dataloader(model.train_dataloader)
        if (self.overfit_batches > 0):