    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"])

        token_loss = nll_loss(
            net_output["token"], sample["target"], sample["mask"], reduce
        )
        dur_loss = self.dur_loss_fn(
            net_output["duration"],
            sample["dur_target"],
            sample["dur_mask"],
            reduce,
        )
        f0_loss = self.f0_loss_fn(
            net_output["f0"],
            sample["f0_target"],
            sample["f0_mask"],
            reduce,
        )
        loss = self.weights.to(token_loss.device) * torch.stack(
            [token_loss, dur_loss, f0_loss], dim=-1
        )
        loss = loss.sum() if reduce else loss.sum(-1)

        sample_size = (
            sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
        )
        logging_output = {
            "loss": loss.detach().sum().item(),
            "token_loss": token_loss.detach().sum().item(),
            "dur_loss": dur_loss.detach().sum().item(),
            "f0_loss": f0_loss.detach().sum().item(),
            "ntokens": sample["ntokens"],
            "nsentences": sample["target"].size(0),
            "sample_size": sample_size,
        }
        return loss, sample_size, 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_ulm_criterion.py