    def _train_step(self, data):
        """Training forward step. Should return average loss of a sample over
        one batch. Provoke ``_forward_step`` is recommended except for special case.
        See ``_train_epoch`` for usage.
        """
        # Init losses
        train_losses = {}
        total_loss = 0

        # Use input feature to get predictions
        data = [b.to(self.accelerator.device, non_blocking=True) for b in data]
        waves, mels, wave_lengths, mel_input_length = data

        # extract semantic latent with w2v model
        waves_16k = torchaudio.functional.resample(waves, 24000, 16000)
        w2v_input = self.w2v_processor(
            waves_16k, sampling_rate=16000, return_tensors="pt"
        ).input_values.to(self.accelerator.device)
        with torch.no_grad():
            w2v_outputs = self.w2v_model(w2v_input.squeeze(0)).logits
            predicted_ids = torch.argmax(w2v_outputs, dim=-1)
            phone_ids = (
                F.interpolate(
                    predicted_ids.unsqueeze(0).float(), mels.size(-1), mode="nearest"
                )
                .long()
                .squeeze(0)
            )

        # get clips
        mel_seg_len = min(
            [int(mel_input_length.min().item()), self.cfg.train.max_frame_len]
        )

        gt_mel_seg = []
        wav_seg = []
        w2v_seg = []

        for bib in range(len(mel_input_length)):
            mel_length = int(mel_input_length[bib].item())

            random_start = (
                np.random.randint(0, mel_length - mel_seg_len)
                if mel_length != mel_seg_len
                else 0
            )
            gt_mel_seg.append(mels[bib, :, random_start : random_start + mel_seg_len])

            # w2v_seg.append(w2v_latent[bib, :, random_start:random_start + mel_seg_len])
            w2v_seg.append(phone_ids[bib, random_start : random_start + mel_seg_len])

            y = waves[bib][random_start * 300 : (random_start + mel_seg_len) * 300]

            wav_seg.append(y.to(self.accelerator.device))

        gt_mel_seg = torch.stack(gt_mel_seg).detach()

        wav_seg = torch.stack(wav_seg).float().detach().unsqueeze(1)
        w2v_seg = torch.stack(w2v_seg).float().detach()

        with torch.no_grad():
            real_norm = log_norm(gt_mel_seg.unsqueeze(1)).squeeze(1).detach()
            F0_real, _, _ = self.pitch_extractor(gt_mel_seg.unsqueeze(1))

        # normalize f0
        # Remove unvoiced frames (replace with -1)
        gt_glob_f0s = []
        f0_targets = []
        for bib in range(len(F0_real)):
            voiced_indices = F0_real[bib] > 5.0
            f0_voiced = F0_real[bib][voiced_indices]

            if len(f0_voiced) != 0:
                # Convert to log scale
                log_f0 = f0_voiced.log2()

                # Calculate mean and standard deviation
                mean_f0 = log_f0.mean()
                std_f0 = log_f0.std()

                # Normalize the F0 sequence
                normalized_f0 = (log_f0 - mean_f0) / std_f0

                # Create the normalized F0 sequence with unvoiced frames
                normalized_sequence = torch.zeros_like(F0_real[bib])
                normalized_sequence[voiced_indices] = normalized_f0
                normalized_sequence[~voiced_indices] = (
                    -10
                )  # Assign -10 to unvoiced frames

                gt_glob_f0s.append(mean_f0)
            else:
                normalized_sequence = torch.zeros_like(F0_real[bib]) - 10.0
                gt_glob_f0s.append(torch.tensor(0.0).to(self.accelerator.device))

            # f0_targets.append(normalized_sequence[single_side_context // 200:-single_side_context // 200])
            f0_targets.append(normalized_sequence)
        f0_targets = torch.stack(f0_targets).to(self.accelerator.device)
        # fill nan with -10
        f0_targets[torch.isnan(f0_targets)] = -10.0
        # fill inf with -10
        f0_targets[torch.isinf(f0_targets)] = -10.0
        # if frame_rate not equal to 80, interpolate f0 from frame rate of 80 to target frame rate
        if self.cfg.preprocess_params.frame_rate != 80:
            f0_targets = F.interpolate(
                f0_targets.unsqueeze(1),
                mel_seg_len // 80 * self.cfg.preprocess_params.frame_rate,
                mode="nearest",
            ).squeeze(1)
            w2v_seg = F.interpolate(
                w2v_seg,
                mel_seg_len // 80 * self.cfg.preprocess_params.frame_rate,
                mode="nearest",
            )

        wav_seg_input = wav_seg
        wav_seg_target = wav_seg

        z = self.model.encoder(wav_seg_input)
        z, quantized, commitment_loss, codebook_loss, timbre = self.model.quantizer(
            z, wav_seg_input, n_c=2, full_waves=waves, wave_lens=wave_lengths
        )
        preds, rev_preds = self.model.fa_predictors(quantized, timbre)

        pred_wave = self.model.decoder(z)

        len_diff = wav_seg_target.size(-1) - pred_wave.size(-1)
        if len_diff > 0:
            wav_seg_target = wav_seg_target[..., len_diff // 2 : -len_diff // 2]

        # discriminator loss
        d_fake = self.model.discriminator(pred_wave.detach())
        d_real = self.model.discriminator(wav_seg_target)
        loss_d = 0
        for x_fake, x_real in zip(d_fake, d_real):
            loss_d += torch.mean(x_fake[-1] ** 2)
            loss_d += torch.mean((1 - x_real[-1]) ** 2)

        self.optimizer.zero_grad()
        self.accelerator.backward(loss_d)
        grad_norm_d = torch.nn.utils.clip_grad_norm_(
            self.model.discriminator.parameters(), 10.0
        )
        self.optimizer.step("discriminator")
        self.optimizer.scheduler(key="discriminator")

        # generator loss
        signal = AudioSignal(wav_seg_target, sample_rate=24000)
        recons = AudioSignal(pred_wave, sample_rate=24000)
        stft_loss = self.criterions["stft"](recons, signal)
        mel_loss = self.criterions["mel"](recons, signal)
        waveform_loss = self.criterions["l1"](recons, signal)

        d_fake = self.model.discriminator(pred_wave)
        d_real = self.model.discriminator(wav_seg_target)

        loss_g = 0
        for x_fake in d_fake:
            loss_g += torch.mean((1 - x_fake[-1]) ** 2)

        loss_feature = 0

        for i in range(len(d_fake)):
            for j in range(len(d_fake[i]) - 1):
                loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())

        pred_f0, pred_uv = preds["f0"], preds["uv"]
        rev_pred_f0, rev_pred_uv = rev_preds["rev_f0"], rev_preds["rev_uv"]

        common_min_size = min(pred_f0.size(-2), f0_targets.size(-1))
        f0_targets = f0_targets[..., :common_min_size]
        real_norm = real_norm[..., :common_min_size]

        f0_loss = F.smooth_l1_loss(
            f0_targets, pred_f0.squeeze(-1)[..., :common_min_size]
        )
        uv_loss = F.smooth_l1_loss(
            real_norm, pred_uv.squeeze(-1)[..., :common_min_size]
        )
        rev_f0_loss = (
            F.smooth_l1_loss(f0_targets, rev_pred_f0.squeeze(-1)[..., :common_min_size])
            if rev_pred_f0 is not None
            else torch.FloatTensor([0]).to(self.accelerator.device)
        )
        rev_uv_loss = (
            F.smooth_l1_loss(real_norm, rev_pred_uv.squeeze(-1)[..., :common_min_size])
            if rev_pred_uv is not None
            else torch.FloatTensor([0]).to(self.accelerator.device)
        )

        tot_f0_loss = f0_loss + rev_f0_loss
        tot_uv_loss = uv_loss + rev_uv_loss

        pred_content = preds["content"]
        rev_pred_content = rev_preds["rev_content"]

        target_content_latents = w2v_seg[..., :common_min_size]

        content_loss = self.criterions["content"](
            pred_content.transpose(1, 2)[..., :common_min_size],
            target_content_latents.long(),
        )
        rev_content_loss = (
            self.criterions["content"](
                rev_pred_content.transpose(1, 2)[..., :common_min_size],
                target_content_latents.long(),
            )
            if rev_pred_content is not None
            else torch.FloatTensor([0]).to(self.accelerator.device)
        )

        tot_content_loss = content_loss + rev_content_loss

        if self.speaker_model is not None:
            spk_logits = torch.cat(
                [
                    self.speaker_model.infer_segment(w16.cpu()[..., :wl])[1]
                    for w16, wl in zip(waves_16k, wave_lengths)
                ],
                dim=0,
            )
            spk_labels = spk_logits.argmax(dim=-1)
        else:
            spk_labels = torch.zeros([len(waves_16k)], dtype=torch.long).to(
                self.accelerator.device
            )

        spk_pred_logits = preds["timbre"]
        spk_loss = F.cross_entropy(spk_pred_logits, spk_labels)
        x_spk_pred_logits = rev_preds["x_timbre"]

        x_spk_loss = (
            F.cross_entropy(x_spk_pred_logits, spk_labels)
            if x_spk_pred_logits is not None
            else torch.FloatTensor([0]).to(self.accelerator.device)
        )

        tot_spk_loss = spk_loss + x_spk_loss

        loss_gen_all = (
            mel_loss * 15.0
            + loss_feature * 1.0
            + loss_g * 1.0
            + commitment_loss * 0.25
            + codebook_loss * 1.0
            + tot_f0_loss * 1.0
            + tot_uv_loss * 1.0
            + tot_content_loss * 5.0
            + tot_spk_loss * 5.0
        )

        self.optimizer.zero_grad()
        self.accelerator.backward(loss_gen_all)

        with torch.no_grad():
            total_loss = loss_gen_all.item()
            train_losses["stft"] = stft_loss.item()
            train_losses["mel"] = mel_loss.item()
            train_losses["l1"] = waveform_loss.item()
            train_losses["f0"] = f0_loss.item()
            train_losses["uv"] = uv_loss.item()
            train_losses["content"] = content_loss.item()
            train_losses["speaker"] = spk_loss.item()
            train_losses["rev_f0"] = rev_f0_loss.item()
            train_losses["rev_uv"] = rev_uv_loss.item()
            train_losses["rev_content"] = rev_content_loss.item()
            train_losses["rev_speaker"] = x_spk_loss.item()

            train_losses["feature"] = loss_feature.item()
            train_losses["generator"] = loss_g.item()
            train_losses["commitment"] = commitment_loss.item()
            train_losses["codebook"] = codebook_loss.item()

            # discriminators
            train_losses["discriminator"] = loss_d.item()

        return total_loss, train_losses
Training forward step. Should return average loss of a sample over
one batch. Provoke ``_forward_step`` is recommended except for special case.
See ``_train_epoch`` for usage.
models/codec/facodec/facodec_trainer.py