    def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        net_output = model(**sample["net_input"])
        loss_dict, stats_dict = self.compute_loss(
            model, net_output, sample, reduce=reduce
        )
        nsentences = sample["net_input"]["src_tokens"][self.channels[0]].size(0)

        logging_output = {
            "nsentences": nsentences,
        }
        logging_output["nsentences"] = nsentences

        loss_all = {t: 0 for t in self.targets}
        correct_all = {t: 0 for t in self.targets}
        count_all = {t: 0 for t in self.targets}
        ntokens_all = 0
        sample_size_all = 0
        for channel in loss_dict:
            for pred_channel in loss_dict[channel]:
                # Get ntokens & sample_size
                ntokens = sample["net_input"]["src_tokens"][channel].numel()
                sample_size = nsentences if self.sentence_avg else ntokens
                prefix = "[{}-{}]".format(channel, pred_channel)
                log_keys = {
                    "next": "general_token",
                    "edge": "edge_token",
                    "duration": "edge_duration",
                }

                # Log & Update the sizes
                logging_output["{}ntokens".format(prefix)] = ntokens
                logging_output["{}sample_size".format(prefix)] = sample_size
                ntokens_all += ntokens
                sample_size_all += sample_size

                for t in self.targets:
                    log_key = log_keys[t]
                    loss = loss_dict[channel][pred_channel][t]
                    correct, count = stats_dict[channel][pred_channel][t]

                    # Log the statistics
                    logging_output["{}{}_loss".format(prefix, log_key)] = loss.data
                    logging_output["{}{}_correct".format(prefix, log_key)] = correct
                    logging_output["{}{}_count".format(prefix, log_key)] = count

                    # Scale the training loss by weights
                    target_loss = loss * self.channel_weights[channel]
                    if pred_channel == channel:
                        target_loss = target_loss * self.main_channel_weight
                    else:
                        target_loss = target_loss * self.cross_channel_weight
                    # Normalize the losses in the training by the number of edges
                    if t in ["edge", "duration"]:
                        target_loss = target_loss / count * sample_size

                    # Update the statistics
                    loss_all[t] += target_loss
                    correct_all[t] += correct
                    count_all[t] += count

        # Logging the average statistics
        logging_output["ntokens"] = ntokens_all
        logging_output["sample_size"] = sample_size_all
        for t in self.targets:
            log_key = {
                "next": "general_token",
                "edge": "edge_token",
                "duration": "edge_duration",
            }[t]
            logging_output["{}_loss".format(log_key)] = loss_all[t].data
            logging_output["{}_correct".format(log_key)] = correct_all[t]
            logging_output["{}_count".format(log_key)] = count_all[t]

        # Define the training loss
        training_loss = 0
        for t in self.targets:
            training_loss += loss_all[t] * self.target_weights[t]
        logging_output["loss"] = training_loss.data

        return training_loss, sample_size_all, logging_output
Compute the loss for the given sample.

Returns a tuple with three elements:
1) the loss
2) the sample size, which is used as the denominator for the gradient
3) logging outputs to display while training
fairseq/criterions/speech_dlm_criterion.py