def _prepare_dataloader(self, dataloader: object, shuffle: bool, mode: RunningStage) -> object:
        """This function handles the following functionalities:

        - Injecting a `DistributedDataSamplerWrapper` into the `DataLoader` if on a distributed environment
        - Wrapping the dataloader based on strategy-specific logic

        """
        # don't do anything if it's not a dataloader
        if not isinstance(dataloader, DataLoader):
            return dataloader
        if (
            self._requires_distributed_sampler(dataloader)  # sets the distributed sampler
            or mode == RunningStage.PREDICTING  # to track indices for the predictions
        ):
            sampler = self._resolve_sampler(dataloader, shuffle=shuffle, mode=mode)
            return _update_dataloader(dataloader, sampler, mode=mode)

        return dataloader
This function handles the following functionalities:

        - Injecting a `DistributedDataSamplerWrapper` into the `DataLoader` if on a distributed environment
        - Wrapping the dataloader based on strategy-specific logic

        
src/lightning/pytorch/trainer/connectors/data_connector.py