def _get_dataset(self, rank: int, num_workers: int, dataloader: _Iterable) -> _List[_Any]:
        """
        Create a subset of dataloader for worker with given rank.
        """
        dataset = []
        num_samples = self._config.calibration_nsamples // num_workers
        sampled = 0
        for idx, data in enumerate(dataloader):
            if idx % num_workers == rank:
                dataset.append(_copy.deepcopy(data))
                sampled += 1
            if sampled == num_samples:
                break
        return dataset

        Create a subset of dataloader for worker with given rank.
        
coremltools/optimize/torch/palettization/sensitive_k_means.py