Update train_ms.py

This commit is contained in:
Stardust·减
2023-08-23 17:06:45 +08:00
committed by GitHub
parent 083500ac06
commit 1fa7a8086b

View File

@@ -331,7 +331,9 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
os.path.join(hps.model_dir, "G_{}.pth".format(global_step)))
utils.save_checkpoint(net_d, optim_d, hps.train.learning_rate, epoch,
os.path.join(hps.model_dir, "D_{}.pth".format(global_step)))
keep_ckpts = getattr(hps.train, 'keep_ckpts', 3)
if net_dur_disc is not None:
utils.save_checkpoint(net_dur_disc, optim_dur_disc, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)))
keep_ckpts = getattr(hps.train, 'keep_ckpts', 5)
if keep_ckpts > 0:
utils.clean_checkpoints(path_to_models=hps.model_dir, n_ckpts_to_keep=keep_ckpts, sort_by_time=True)