    def _train_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]
        """
        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,T] -> (n_q, B, T)
                # RVQ_1 = codes[:1, :, :] # Contain content info, can be considered as semantic tokens
                # RVQ_supplement = codes[1:, :, :] # Contain timbre info, complete info lost by the first quantizer
                # 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)

                # # Decoding from RVQ-i:j tokens from the ith quantizers to the jth quantizers
                # wav = model.decode(codes[i: (j + 1)], st=i)
            else:
                # using encodec, 24k
                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//320]
            batch["speech"] = vq_id
        batch["speech_len"] = batch["speech_len"] // 320  # our codec downsamples 320x
        assert batch["speech_len"].max() <= batch["speech"].shape[-1]

        phone_mask = 1 - make_pad_mask(
            batch["phone_lens"], max_len=batch["phone_ids"].size(1), left_pad=False
        ).to(torch.long)
        speech_mask = 1 - make_pad_mask(
            batch["speech_len"], max_len=batch["speech"].size(-1)
        ).to(torch.long)

        np.random.seed(int(time.time()) - 5 * self.accelerator.process_index)

        if hasattr(self.cfg.train, "dropout"):
            dropout = self.cfg.train.dropout
        else:
            dropout = 0.0

        out = self.model(
            phone_ids=batch["phone_ids"],
            phone_mask=phone_mask,
            target_ids=batch["speech"],
            target_mask=speech_mask,
            dropout=dropout,
        )
        loss = out.loss

        self.accelerator.log(
            {f"Train/NAR L{out.target_quantization_layer} Top1 acc": out.top1_acc},
            step=self.step,
        )
        self.accelerator.log(
            {f"Train/NAR L{out.target_quantization_layer} Top5 acc": out.top5_acc},
            step=self.step,
        )
        self.accelerator.log(
            {f"Train/NAR L{out.target_quantization_layer} Top10 acc": out.top10_acc},
            step=self.step,
        )

        # if hasattr(out, 'top1_acc'):
        #     idx = out.target_quantization_layer
        #     self.top1_accuracies[idx].append(out.top1_acc)
        #     self.top5_accuracies[idx].append(out.top5_acc)
        #     self.top10_accuracies[idx].append(out.top10_acc)
        #     if len(self.top1_accuracies[idx]) >= 160:
        #         breakpoint()
        # if self.accelerator.is_main_process:
        #     print(loss)
        return loss
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