Update train_ms.py

This commit is contained in:
Stardust·减
2023-08-21 17:17:40 +08:00
committed by GitHub
parent 274e3e23b2
commit 12c9fb9ee8

View File

@@ -133,13 +133,13 @@ def run(rank, n_gpus, hps):
if pretrain_dir is None: if pretrain_dir is None:
try: try:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
optim_g, skip_optimizer) optim_g, skip_optimizer=True)
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
optim_d, skip_optimizer) optim_d, skip_optimizer=True)
epoch_str = max(epoch_str, 1) epoch_str = max(epoch_str, 1)
global_step = (epoch_str - 1) * len(train_loader) global_step = (epoch_str - 1) * len(train_loader)
except Exception: except Exception as e:
print("load old checkpoint failed...") print(e)
epoch_str = 1 epoch_str = 1
global_step = 0 global_step = 0
else: else: