def prepare_data_loader(self,
                            data_loader: torch.utils.data.DataLoader,
                            add_dist_sampler: bool = True,
                            move_to_device: bool = True) -> torch.utils.data.DataLoader:
        """
        Prepares DataLoader for distributed execution.

        This allows you to use the same exact code regardless of number of
        workers or the device type being used (CPU, GPU).

        Args:
            data_loader (torch.utils.data.DataLoader): The DataLoader to
                prepare.
            add_dist_sampler (bool): Whether to add a DistributedSampler to
                the provided DataLoader.
            move_to_device (bool): If set, automatically move the data
                returned by the data loader to the correct device.
        """

        # Only add Distributed Sampler if the following conditions hold:
        # 1. More than one training worker is being used.
        # 2. A DistributedSampler has not already been added by the user.
        # 3. The dataset is not an IterableDataset. Samplers do not worker with
        # IterableDatasets.
        def with_sampler(loader):
            # Automatically set the DistributedSampler

            # If using a sampler, the shuffle attribute in the
            # DataLoader must be set to False.
            # Instead the shuffling is determined by the shuffle attribute
            # in the DistributedSampler.
            # We identify if shuffling is enabled in the passed in
            # DataLoader by seeing if the sampler for the DataLoader is a
            # SequentialSampler.
            shuffle = not isinstance(loader.sampler, SequentialSampler)

            data_loader_args = {
                "dataset": loader.dataset,
                "batch_size": loader.batch_size,
                "shuffle": False,
                "num_workers": loader.num_workers,
                "collate_fn": loader.collate_fn,
                "pin_memory": loader.pin_memory,
                "drop_last": loader.drop_last,
                "timeout": loader.timeout,
                "worker_init_fn": loader.worker_init_fn,
                "sampler": DistributedSampler(loader.dataset, shuffle=shuffle)
            }
            return DataLoader(**data_loader_args)

        data_loader = with_sampler(data_loader)

        if move_to_device:
            if torch.cuda.is_available():
                rank = self.local_rank
                device = torch.device(f"cuda:{rank}")
            else:
                device = torch.device("cpu")
            data_loader = DistributedDL(data_loader, device)

        return data_loader

        Prepares DataLoader for distributed execution.

        This allows you to use the same exact code regardless of number of
        workers or the device type being used (CPU, GPU).

        Args:
            data_loader (torch.utils.data.DataLoader): The DataLoader to
                prepare.
            add_dist_sampler (bool): Whether to add a DistributedSampler to
                the provided DataLoader.
            move_to_device (bool): If set, automatically move the data
                returned by the data loader to the correct device.
        
PatchTST_self_supervised/src/callback/distributed.py