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, codebook_out, encoder_output = model(**sample["net_input"])
        bart_loss, _ = self.compute_loss(model, net_output, sample, reduce=reduce)
        sample_size = (
            sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
        )
        
        loss = self.bart_weight * bart_loss
        logging_output = {
            "loss": loss.item(),
            "ntokens": sample["ntokens"],
            "nsentences": sample["target"].size(0),
            "bart_loss": bart_loss.item(),
            "sample_size": sample_size,
        }

        if "prob_perplexity" in codebook_out:
            assert hasattr(model, "get_extra_losses")
            extra_losses, names = model.get_extra_losses(codebook_out)
            if torch.is_tensor(extra_losses):
                extra_losses = [extra_losses]
                names = [names]
            if len(self.loss_weights) == 1 and len(extra_losses) != 1:
                self.loss_weights = [self.loss_weights[0]] * len(extra_losses)
            if len(self.loss_weights) > len(extra_losses):
                modified_loss_weight = self.loss_weights[len(extra_losses):]
            else:
                modified_loss_weight = self.loss_weights

            # assert len(extra_losses) == len(self.loss_weights), f"{len(extra_losses)}, {len(self.loss_weights)}"
            for p, n, coef in zip(extra_losses, names, modified_loss_weight):
                # print(n + str(coef))
                if coef != 0 and p is not None:
                    p = coef * p.float() * sample_size
                    loss += p
                    logging_output[f"loss_{n}"] = p.item()

        if 'loss_prob_perplexity' in logging_output:
            logging_output['code_perplexity'] = codebook_out['code_perplexity'].item()

        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
        
speecht5/speecht5/criterions/text_pretrain_criterion.py