    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
                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
                )
            # 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])

            if self.flatten_first_2_layers:
                first_layer = vq_id[0]
                second_layer = vq_id[1]
                # flatten the first two layers
                batch["speech"] = torch.stack(
                    [first_layer, second_layer], dim=-1
                ).flatten(-2, -1)
                batch["speech_len"] = batch["speech_len"] // 160
            elif hasattr(self.cfg.model, "num_prediction_heads"):
                batch["speech"] = vq_id[:2]  # first two layers
                batch["speech_len"] = (
                    batch["speech_len"] // 320
                )  # our codec downsamples 320x
            else:
                batch["speech"] = vq_id[0]  # use first layer
                batch["speech_len"] = (
                    batch["speech_len"] // 320
                )  # our codec downsamples 320x

            # save gt
            breakpoint()
            recovered_audio = self.codec_encoder.decode(vq_id[:1, :1])
            # recovered_audio = self.codec_encoder.decode([(vq_id[:1].transpose(0,1), None)])
            torchaudio.save("gt.wav", recovered_audio[0].cpu(), 16000)
            out_vq_ids = self.model.sample_hf(
                batch["phone_ids"][:1, ...], batch["speech"][:1, :225], temperature=0.9
            )
            # out_vq_ids = torch.cat([batch['speech'][:1, :225], out_vq_ids[:1, ...]], dim=1)

            # reconstruct form tokens
            recovered_audio = self.codec_encoder.decode(out_vq_ids.unsqueeze(0))
            # recovered_audio = self.codec_encoder.decode([(out_vq_ids, None)])
            torchaudio.save("a.wav", recovered_audio[0].cpu(), 16000)
            breakpoint()
            print()
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_ar_trainer.py