    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
                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//320]
            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
        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)

        out = self.model(
            phone_ids=batch["phone_ids"],
            phone_mask=phone_mask,
            target_ids=batch["speech"],
            target_mask=speech_mask,
        )
        loss = out.loss
        # if self.accelerator.is_main_process:
        #     print(loss)
        # if hasattr(out, 'top1_acc'):
        #     self.top1_accuracies.append(out.top1_acc)
        #     self.top5_accuracies.append(out.top5_acc)
        #     self.top10_accuracies.append(out.top10_acc)
        #     print(f'avgs: top1: {sum(self.top1_accuracies)/len(self.top1_accuracies)}, top5: {sum(self.top5_accuracies)/len(self.top5_accuracies)}, top10: {sum(self.top10_accuracies)/len(self.top10_accuracies)}')
        #     breakpoint()
        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_ar_trainer.py