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

        sample_size = sample["id"].numel()
        assert sample_size % self.task.cfg.mt_beam == 0, (
            f"sample_size ({sample_size}) cannot be divided by beam size ({self.task.cfg.mt_beam})."
            f"Please set --required-batch-size-multiple={self.task.cfg.mt_beam}."
        )

        # split into smaller batches for model forward
        batch_out = []
        for i in range(0, sample_size, self.forward_batch_size):
            j = min(i + self.forward_batch_size, sample_size)

            out = model(
                src_tokens=sample["net_input"]["src_tokens"][i:j, :],
                src_lengths=sample["net_input"]["src_lengths"][i:j],
            )

            batch_out.append(
                model.sentence_forward(out, sample["net_input"]["src_tokens"][i:j, :])
            )

        batch_out = torch.cat(batch_out, dim=0).view(
            self.task.cfg.mt_beam, sample_size // self.task.cfg.mt_beam, -1
        )  # T x B x C
        if model.joint_classification == "sent":
            batch_out = model.joint_forward(batch_out)
        scores = model.classification_forward(batch_out.view(sample_size, 1, -1)).view(
            -1, self.task.cfg.mt_beam
        )  # input: B x T x C

        loss = self.compute_kl_loss(
            scores, sample["target"][:, 0].view(-1, self.task.cfg.mt_beam)
        )

        sample_size = sample_size // self.task.cfg.mt_beam

        logging_output = {
            "loss": loss.detach(),
            "ntokens": sample["ntokens"],
            "nsentences": sample_size * self.task.cfg.mt_beam,
            "sample_size": sample_size,
            "scores": scores.detach(),
        }

        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
examples/discriminative_reranking_nmt/criterions/discriminative_reranking_criterion.py