diff --git a/train_ms.py b/train_ms.py
index 4f0cd69..ddf6506 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -155,34 +155,20 @@ def run(rank, n_gpus, hps):
     net_d = DDP(net_d, device_ids=[rank], find_unused_parameters=True)
     if net_dur_disc is not None:
         net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)
-
-    pretrain_dir = None
-    if pretrain_dir is None:
-        try:
-            if net_dur_disc is not None:
-                _, optim_dur_disc, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, skip_optimizer=not hps.resume)
-                _, optim_g, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
-                                                   optim_g, skip_optimizer=not hps.resume)
-                _, optim_d, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
-                                                   optim_d, skip_optimizer=not hps.resume)
-            else:
-                _, optim_g, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
-                                                   optim_g, skip_optimizer=not hps.resume)
-                _, optim_d, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
-                                                   optim_d, skip_optimizer=not hps.resume)
+    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)
+        _, optim_g, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
+                                                   optim_g, skip_optimizer=True)
+        _, optim_d, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
+                                                   optim_d, skip_optimizer=True)
             
-            epoch_str = max(epoch_str, 1)
-            global_step = (epoch_str - 1) * len(train_loader)
+        epoch_str = max(epoch_str, 1)
+        global_step = (epoch_str - 1) * len(train_loader)
         except Exception as e:
             print(e)
             epoch_str = 1
             global_step = 0
-    else:
-        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(pretrain_dir, "G_*.pth"), net_g,
-                                                   optim_g, True)
-        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(pretrain_dir, "D_*.pth"), net_d,
-                                                   optim_d, True)
-
 
 
     scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)