From 51a03dbc17b525008496c28aea0ed953a8b00195 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Wed, 23 Aug 2023 19:00:21 +0800
Subject: [PATCH] Update train_ms.py

---
 train_ms.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/train_ms.py b/train_ms.py
index a18892a..e2c7149 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -156,12 +156,13 @@ def run(rank, n_gpus, hps):
     pretrain_dir = None
     if pretrain_dir is None:
         try:
+            if net_dur_disc is not None:
+                _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, skip_optimizer=True)
             _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
                                                    optim_g, skip_optimizer=True)
             _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
                                                    optim_d, skip_optimizer=True)
-            #if net_dur_disc is not None:
-               # _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, skip_optimizer=True)
+            
             epoch_str = max(epoch_str, 1)
             global_step = (epoch_str - 1) * len(train_loader)
         except Exception as e: