diff --git a/train_ms.py b/train_ms.py
index 54b1819..6b59fd2 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -98,7 +98,7 @@ def run(rank, n_gpus, hps):
     if "use_duration_discriminator" in hps.model.keys() and hps.model.use_duration_discriminator == True:
         print("Using duration discriminator for VITS2")
         use_duration_discriminator = True
-       net_dur_disc = DurationDiscriminator(
+        net_dur_disc = DurationDiscriminator(
          hps.model.hidden_channels, 
          hps.model.hidden_channels, 
          3, 
@@ -265,7 +265,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
             # Generator
             y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
             if net_dur_disc is not None:
-                y_dur_hat_r, y_dur_hat_g = net_dur_disc(logw, logw_)
+                y_dur_hat_r, y_dur_hat_g = net_dur_disc(x, x_mask, logw, logw_)
             with autocast(enabled=False):
                 loss_dur = torch.sum(l_length.float())
                 loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
@@ -274,6 +274,9 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
                 loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+                if net_dur_disc is not None:
+                    loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
+                    loss_gen_all += loss_dur_gen
         optim_g.zero_grad()
         scaler.scale(loss_gen_all).backward()
         scaler.unscale_(optim_g)
@@ -302,6 +305,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
                 if net_dur_disc is not None:
                     scalar_dict.update({f"loss/dur_disc_r/{losses_dur_disc_r}"})
                     scalar_dict.update({f"loss/dur_disc_g/{losses_dur_disc_g}"})
+                    scalar_dict.update({f"loss/dur_gen/{loss_dur_gen}"})
                 image_dict = {
                     "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
                     "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),