    def _test_step(self, batch):
        # inference codec
        """Returns: dict('speech', 'speech_len', 'phone_ids', 'phone_lens')
        speech: [B, T]
        speech_len: [B]
        phone_ids: [B, T]
        phone_lens: [B]
        """
        import torchaudio

        device = self.accelerator.device
        for k, v in batch.items():
            if isinstance(v, torch.Tensor):
                batch[k] = v.to(device)
        with torch.no_grad():
            if self.cfg.use_speechtokenizer:
                # Extract discrete codes from SpeechTokenizer
                # 16k
                vq_id = self.codec_encoder.encode(
                    batch["speech"].unsqueeze(1)
                )  # [B,1,T] -> (n_q, B, T)
                # Concatenating semantic tokens (RVQ_1) and supplementary timbre tokens and then decoding
                # wav = self.codec_encoder.decode(vq_id)
                # torchaudio.save('a.wav', wav[0].cpu(), 16000)

            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
                )
            # recovered_audio = self.codec_encoder.decode([(vq_id.transpose(0,1), None)])
            # recovered_audio = self.codec_decoder(vq_emb, vq=False)
            # torchaudio.save('a.wav', recovered_audio[0], 16000)
            # vq_id: [8, B, T//200]

            # vq_emb = self.codec_decoder.quantizer.vq2emb(vq=vq_id[:1], n_quantizers=1)
            # recovered_audio = self.codec_decoder(vq_emb, vq=False)
            # recovered_audio.shape: torch.Size([1, 1, 50200])

            batch["speech"] = vq_id

            # save gt
            if self.cfg.use_speechtokenizer:
                recovered_audio = self.codec_encoder.decode(vq_id)
            else:
                recovered_audio = self.codec_encoder.decode(
                    [(vq_id.transpose(0, 1), None)]
                )
            torchaudio.save("gt.wav", recovered_audio[0].cpu(), 16000)
            self.model.eval()
            out_vq_ids = self.model.sample_hf(
                phone_ids=batch["phone_ids"][:1],
                prompt_ids=batch["speech"][:, :1, :150],
                first_stage_ids=batch["speech"][0, :1, 150:],
            )
            # breakpoint()
            # out_vq_ids = torch.cat([batch['speech'][:, :225], out_vq_ids], dim=1)

            # reconstruct form tokens
            if self.cfg.use_speechtokenizer:
                recovered_audio = self.codec_encoder.decode(out_vq_ids)
            else:
                recovered_audio = self.codec_encoder.decode(
                    [(out_vq_ids.transpose(0, 1)[:1], None)]
                )
            torchaudio.save("a.wav", recovered_audio[0].cpu(), 16000)
            breakpoint()
Returns: dict('speech', 'speech_len', 'phone_ids', 'phone_lens')
speech: [B, T]
speech_len: [B]
phone_ids: [B, T]
phone_lens: [B]
models/tts/valle_v2/valle_nar_trainer.py