    def _build_criterion(self):
        class GeneratorLoss(nn.Module):
            def __init__(self, cfg):
                super(GeneratorLoss, self).__init__()
                self.cfg = cfg
                self.l1_loss = nn.L1Loss()

            def generator_loss(self, disc_outputs):
                loss = 0
                gen_losses = []
                for dg in disc_outputs:
                    dg = dg.float()
                    l = torch.mean((1 - dg) ** 2)
                    gen_losses.append(l)
                    loss += l

                return loss, gen_losses

            def feature_loss(self, fmap_r, fmap_g):
                loss = 0
                for dr, dg in zip(fmap_r, fmap_g):
                    for rl, gl in zip(dr, dg):
                        rl = rl.float().detach()
                        gl = gl.float()
                        loss += torch.mean(torch.abs(rl - gl))

                return loss * 2

            def kl_loss(self, z_p, logs_q, m_p, logs_p, z_mask):
                """
                z_p, logs_q: [b, h, t_t]
                m_p, logs_p: [b, h, t_t]
                """
                z_p = z_p.float()
                logs_q = logs_q.float()
                m_p = m_p.float()
                logs_p = logs_p.float()
                z_mask = z_mask.float()

                kl = logs_p - logs_q - 0.5
                kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p)
                kl = torch.sum(kl * z_mask)
                l = kl / torch.sum(z_mask)
                return l

            def forward(
                self,
                outputs_g,
                outputs_d,
                y_mel,
                y_hat_mel,
            ):
                loss_g = {}

                # mel loss
                loss_mel = self.l1_loss(y_mel, y_hat_mel) * self.cfg.train.c_mel
                loss_g["loss_mel"] = loss_mel

                # kl loss
                loss_kl = (
                    self.kl_loss(
                        outputs_g["z_p"],
                        outputs_g["logs_q"],
                        outputs_g["m_p"],
                        outputs_g["logs_p"],
                        outputs_g["z_mask"],
                    )
                    * self.cfg.train.c_kl
                )
                loss_g["loss_kl"] = loss_kl

                # feature loss
                loss_fm = self.feature_loss(outputs_d["fmap_rs"], outputs_d["fmap_gs"])
                loss_g["loss_fm"] = loss_fm

                # gan loss
                loss_gen, losses_gen = self.generator_loss(outputs_d["y_d_hat_g"])
                loss_g["loss_gen"] = loss_gen
                loss_g["loss_gen_all"] = loss_mel + loss_kl + loss_fm + loss_gen

                return loss_g

        class DiscriminatorLoss(nn.Module):
            def __init__(self, cfg):
                super(DiscriminatorLoss, self).__init__()
                self.cfg = cfg
                self.l1Loss = torch.nn.L1Loss(reduction="mean")

            def __call__(self, disc_real_outputs, disc_generated_outputs):
                loss_d = {}

                loss = 0
                r_losses = []
                g_losses = []
                for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
                    dr = dr.float()
                    dg = dg.float()
                    r_loss = torch.mean((1 - dr) ** 2)
                    g_loss = torch.mean(dg**2)
                    loss += r_loss + g_loss
                    r_losses.append(r_loss.item())
                    g_losses.append(g_loss.item())

                loss_d["loss_disc_all"] = loss

                return loss_d

        criterion = {
            "generator": GeneratorLoss(self.cfg),
            "discriminator": DiscriminatorLoss(self.cfg),
        }
        return criterion
z_p, logs_q: [b, h, t_t]
m_p, logs_p: [b, h, t_t]
models/svc/vits/vits_trainer.py