    def forward(
        self,
        x: torch.Tensor,
        x_lens: torch.Tensor,
        y: Union[torch.Tensor, PromptedFeatures],
        y_lens: Union[torch.Tensor, PromptedFeatures],
        reduction: str = "sum",
        train_stage: int = 0,
        **kwargs,
    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
        """
        Args:
          x:
            A 2-D tensor of shape (N, S).
          x_lens:
            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
            before padding.
          y:
            A 3-D tensor of shape (N, T, 8).
          y_lens:
            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
            before padding.
          train_stage:
            0: AR & NAR modules, 1: AR modules, 2: NAR modules
        Returns:
          Return the predicted audio code matrix, cross-entropy loss and Top-10 accuracy.
        """
        assert x.ndim == 2, x.shape
        assert x_lens.ndim == 1, x_lens.shape

        y_prompts_codes = None
        if isinstance(y, PromptedFeatures):
            y_prompts_codes, y = y.data
            prompts_len, y_lens = y_lens.data
            assert prompts_len.min() == prompts_len.max()
            assert self.prefix_mode == 4
            y_prompts_codes = y_prompts_codes.type(torch.int64)

        assert y.ndim == 3, y.shape
        assert y_lens.ndim == 1, y_lens.shape

        x_mask = make_pad_mask(x_lens).to(x.device)
        y_mask = make_pad_mask(y_lens).to(y.device)
        y_mask_int = y_mask.type(torch.int64)

        text = x
        codes = y.type(torch.int64) * (1 - y_mask_int.unsqueeze(dim=-1))

        y, targets = self.pad_y_eos(
            codes[..., 0], y_mask_int, eos_id=self.audio_token_num
        )
        self.y_mask_int = y_mask_int

        metrics = {}
        total_loss = 0.0

        xy_padding_mask = torch.concat([x_mask, y_mask], dim=1)
        if self.ar_audio_prepend_bos:
            ar_xy_padding_mask = torch.concat(
                [x_mask, F.pad(y_mask, (1, 0), value=False)], dim=1
            )
        else:
            ar_xy_padding_mask = xy_padding_mask
        self.xy_padding_mask = xy_padding_mask
        self.ar_xy_padding_mask = ar_xy_padding_mask

        # AR Decoder
        if train_stage in [0, 1]:
            ar_loss, ar_metrics = self._forward_ar_decoder(
                text, x_lens.max(), y, y_lens.max(), targets, x_mask, y_mask, reduction
            )
            total_loss += ar_loss
            metrics["AR_Top100Acc"] = ar_metrics

        # NAR Decoder
        if self.ar_audio_prepend_bos:
            y = y[:, 1:]

        if self.num_quantizers > 1 and train_stage in [0, 2]:
            nar_loss, nar_metrics = self._forward_nar_decoder(
                text,
                x_lens,
                y,
                y_lens,
                codes,
                y_prompts_codes,
                x_mask,
                y_mask,
                reduction,
            )
            total_loss += nar_loss
            metrics["NAR_Top100Acc"] = nar_metrics

        if train_stage == 0:
            total_loss = total_loss / 2.0

        return total_loss, metrics
Args:
x:
A 2-D tensor of shape (N, S).
x_lens:
A 1-D tensor of shape (N,). It contains the number of tokens in `x`
before padding.
y:
A 3-D tensor of shape (N, T, 8).
y_lens:
A 1-D tensor of shape (N,). It contains the number of tokens in `x`
before padding.
train_stage:
0: AR & NAR modules, 1: AR modules, 2: NAR modules
Returns:
Return the predicted audio code matrix, cross-entropy loss and Top-10 accuracy.
models/tts/valle/valle.py