def collater(self, samples):
        """
        If we are doing batch sampling, then pick the right collater to use.

        Otherwise we assume all collaters are the same.
        """
        if len(samples) == 0:
            return None
        
        samples_dict = {key: [] for key in self.datasets}
        for s in samples:
            _, key = self._map_index(s["full_id"])
            samples_dict[key].append(s)
        
        batch = {}
        for key in samples_dict:
            if len(samples_dict[key]) == 0:
                continue
            batch[key] = self.datasets[key].collater(samples_dict[key])

        return batch

        If we are doing batch sampling, then pick the right collater to use.

        Otherwise we assume all collaters are the same.
        
speechlm/speechlm/data/multimodal_corpus_dataset.py