    def forward(self, batch, chunk_configs: list, return_prompt=False, prompt_len=None):
        """batch: dict(
            speech: [B, T]
            phone_ids: [B, T]
        )
        returns: [B, 1, T] audio
        """
        if prompt_len is None:
            prompt_len = 100000  # no prompt length limiting
        for k, v in batch.items():
            if isinstance(v, torch.Tensor):
                batch[k] = v.to(self.device)
        with torch.no_grad():
            if self.use_speechtokenizer:
                vq_id = self.codec_encoder.encode(
                    batch["speech"].unsqueeze(1)
                )  # [B,1,T] -> (n_q, B, T)
            else:
                vq_id = self.codec_encoder.encode(batch["speech"].unsqueeze(1))
                vq_id = torch.cat([encoded[0] for encoded in vq_id], dim=-1).transpose(
                    0, 1
                )

            # typically we only require one config in the chunk,
            # but we can also use multiple configs to, for example, use different sampling temperature at different positions
            for chunk in chunk_configs:
                ar_vq_ids = self.ar_model.sample_hf(
                    batch["phone_ids"],
                    vq_id[0, :, :prompt_len],
                    top_p=chunk["top_p"],
                    top_k=chunk["top_k"],
                    temperature=chunk["temperature"],
                    num_beams=chunk["num_beams"],
                    repeat_penalty=chunk["repeat_penalty"],
                    max_length=chunk["max_length"],
                )
                # recovered_audio_ar = self.decode(ar_vq_ids.unsqueeze(0))
                # torchaudio.save('recovered_audio_ar.wav', recovered_audio_ar[0].cpu(), 24000)

                nar_vq_ids = self.nar_model.sample_hf(
                    phone_ids=batch["phone_ids"],
                    prompt_ids=vq_id[:, :, :prompt_len],
                    first_stage_ids=ar_vq_ids,
                    # first_stage_ids=vq_id[0, :, prompt_len:],
                )

                if return_prompt:
                    nar_vq_ids = torch.cat(
                        [vq_id[..., :prompt_len], nar_vq_ids], dim=-1
                    )

                recovered_audio = self.decode(nar_vq_ids)
                return recovered_audio  # [B, 1, T]
batch: dict(
speech: [B, T]
phone_ids: [B, T]
)
returns: [B, 1, T] audio
models/tts/valle_v2/valle_inference.py