    def CTLoss_D(self, y, cond, mask):
        """
        compute loss for consistency distillation

        Args:
            y: ground truth mel-spectrogram [B x n_mel x L]
            cond: output of conformer encoder [B x n_mel x L]
            mask: mask of padded frames [B x n_mel x L]
        """
        with torch.no_grad():
            mu = 0.95
            for p, ema_p in zip(
                self.denoise_fn.parameters(), self.denoise_fn_ema.parameters()
            ):
                ema_p.mul_(mu).add_(p, alpha=1 - mu)

        n = torch.randint(1, self.N, (y.shape[0],))
        z = torch.randn_like(y) + cond

        tn_1 = self.t_steps[n + 1].reshape(-1, 1, 1).to(y.device)
        f_theta = self.EDMPrecond(y + tn_1 * z, tn_1, cond, self.denoise_fn)

        with torch.no_grad():
            tn = self.t_steps[n].reshape(-1, 1, 1).to(y.device)

            # euler step
            x_hat = y + tn_1 * z
            denoised = self.EDMPrecond(x_hat, tn_1, cond, self.denoise_fn_pretrained)
            d_cur = (x_hat - denoised) / tn_1
            y_tn = x_hat + (tn - tn_1) * d_cur

            # Heun’s 2nd order method

            denoised2 = self.EDMPrecond(y_tn, tn, cond, self.denoise_fn_pretrained)
            d_prime = (y_tn - denoised2) / tn
            y_tn = x_hat + (tn - tn_1) * (0.5 * d_cur + 0.5 * d_prime)

            f_theta_ema = self.EDMPrecond(y_tn, tn, cond, self.denoise_fn_ema)

        loss = (f_theta - f_theta_ema.detach()) ** 2
        loss = torch.sum(loss * mask) / torch.sum(mask)

        # check nan
        if torch.any(torch.isnan(loss)):
            print("nan loss")
        if torch.any(torch.isnan(f_theta)):
            print("nan f_theta")
        if torch.any(torch.isnan(f_theta_ema)):
            print("nan f_theta_ema")

        return loss
compute loss for consistency distillation

Args:
y: ground truth mel-spectrogram [B x n_mel x L]
cond: output of conformer encoder [B x n_mel x L]
mask: mask of padded frames [B x n_mel x L]
models/svc/comosvc/comosvc.py