    def collater(self, samples):
        """Merge a list of samples to form a mini-batch.

        Args:
            samples (List[dict]): samples to collate

        Returns:
            dict: a mini-batch with the following keys:

                - `id` (LongTensor): example IDs in the original input order
                - `ntokens` (int): total number of tokens in the batch
                - `net_input` (dict): the input to the Model, containing keys:

                  - `src_tokens` (OrderedDict[str, LongTensor]): dictionary
                    over channel with the values being padded 2D Tensor of
                    samples `source` of shape `(bsz, src_len)`.
                    Padding will appear on the right.
                  - `src_lengths` (LongTensor): lengths of source sentences
                    in the mini-batch

                - `target` (dict): the target of the Model, containing keys:

                  - `next` (OrderedDict[str, LongTensor]): dictionary
                    over channel with the values being padded 2D Tensor of
                    batch samples' `target_next` of shape `(bsz, tgt_len)`.
                    Padding will appear on the right.
                  - `edge` (OrderedDict[str, LongTensor]): dictionary
                    over channel with the values being the concatenated
                    1D Tensor of batch samples' `target_edge` of shape
                    `(sum of dedup_tgt_len,)`
                  - `duration` (OrderedDict[str, LongTensor]): dictionary
                    over channel with the values being the concatenated
                    1D Tensor of batch samples' `target_duration` of shape
                    `(sum of dedup_tgt_len,)`
                  - `edge_indices` (OrderedDict[str, LongTensor]): dictionary
                    over channel with the values being the concatenated
                    1D Tensor of batch samples' `target_edge_indices` of
                    shape `(sum of dedup_tgt_len,)`.
                    The indices are added to multiplies of batch size
                    such that they are the actual indices in the flatten
                    `src_tokens` Tensor
        """
        if len(samples) == 0:
            return {}

        pad_idx = self.vocab.pad()
        eos_idx = self.vocab.eos()

        def merge(key, max_size=None):
            if samples[0][key] is None:
                return None
            res = OrderedDict()
            for channel in samples[0][key]:
                if key in ["source", "target_next"]:
                    # fill batch of shape: (batch_size, max_size)
                    res[channel] = data_utils.collate_tokens(
                        [s[key][channel] for s in samples],
                        pad_idx,
                        eos_idx,
                        left_pad=False,
                    )
                elif key in ["target_edge", "target_duration"]:
                    # concatenate the edge units/duration
                    res[channel] = torch.cat([s[key][channel] for s in samples])
                elif key == "target_edge_indices":
                    # increase the edge indices to the indices in the flatten batch
                    res[channel] = torch.cat(
                        [s[key][channel] + i * max_size for i, s in enumerate(samples)]
                    )

            return res

        src_tokens = merge("source")
        tgt_next = merge("target_next")
        tgt_edge = merge("target_edge")
        tgt_duration = merge("target_duration")
        tgt_edge_indices = merge(
            "target_edge_indices", max_size=next(iter(src_tokens.values())).size(-1)
        )
        return {
            "id": torch.LongTensor([s["id"] for s in samples]),
            "nsentences": len(samples),
            "ntokens": sum(len(item) for s in samples for item in s["source"].values()),
            "net_input": {
                "src_tokens": src_tokens,
                "src_lengths": torch.LongTensor(
                    [next(iter(s["source"].values())).numel() for s in samples]
                ),
            },
            "target": {
                "next": tgt_next,
                "edge": tgt_edge,
                "duration": tgt_duration,
                "edge_indices": tgt_edge_indices,
            },
        }
Merge a list of samples to form a mini-batch.

Args:
samples (List[dict]): samples to collate

Returns:
dict: a mini-batch with the following keys:

- `id` (LongTensor): example IDs in the original input order
- `ntokens` (int): total number of tokens in the batch
- `net_input` (dict): the input to the Model, containing keys:

- `src_tokens` (OrderedDict[str, LongTensor]): dictionary
over channel with the values being padded 2D Tensor of
samples `source` of shape `(bsz, src_len)`.
Padding will appear on the right.
- `src_lengths` (LongTensor): lengths of source sentences
in the mini-batch

- `target` (dict): the target of the Model, containing keys:

- `next` (OrderedDict[str, LongTensor]): dictionary
over channel with the values being padded 2D Tensor of
batch samples' `target_next` of shape `(bsz, tgt_len)`.
Padding will appear on the right.
- `edge` (OrderedDict[str, LongTensor]): dictionary
over channel with the values being the concatenated
1D Tensor of batch samples' `target_edge` of shape
`(sum of dedup_tgt_len,)`
- `duration` (OrderedDict[str, LongTensor]): dictionary
over channel with the values being the concatenated
1D Tensor of batch samples' `target_duration` of shape
`(sum of dedup_tgt_len,)`
- `edge_indices` (OrderedDict[str, LongTensor]): dictionary
over channel with the values being the concatenated
1D Tensor of batch samples' `target_edge_indices` of
shape `(sum of dedup_tgt_len,)`.
The indices are added to multiplies of batch size
such that they are the actual indices in the flatten
`src_tokens` Tensor
fairseq/data/speech_dlm_dataset.py