    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
        total_loss = 0

        # Use input feature to get predictions
        mel_input = data["mel"]
        audio_gt = data["audio"]

        if self.cfg.preprocess.use_frame_pitch:
            pitch_input = data["frame_pitch"]

        self.optimizer.zero_grad()
        N = audio_gt.shape[0]
        t = torch.randint(
            0, len(self.cfg.model.diffwave.noise_schedule), [N], device=self.device
        )
        noise_scale = self.noise_level[t].unsqueeze(1)
        noise_scale_sqrt = noise_scale**0.5
        noise = torch.randn_like(audio_gt).to(self.device)
        noisy_audio = noise_scale_sqrt * audio_gt + (1.0 - noise_scale) ** 0.5 * noise

        audio_pred = self.model(noisy_audio, t, mel_input)
        total_loss = self.criterion(noise, audio_pred.squeeze(1))

        self.accelerator.backward(total_loss)
        self.optimizer.step()

        return total_loss.item()
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/vocoders/diffusion/diffusion_vocoder_trainer.py