From ddfede1f52191fd9568a9fe5c178b2147b6e0b3e Mon Sep 17 00:00:00 2001 From: erogol Date: Wed, 10 Jun 2020 13:55:37 +0200 Subject: [PATCH] restore vocoder scheduler if it is in checkpoint --- vocoder/train.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/vocoder/train.py b/vocoder/train.py index a8a8f011..091a6932 100644 --- a/vocoder/train.py +++ b/vocoder/train.py @@ -450,10 +450,14 @@ def main(args): # pylint: disable=redefined-outer-name weight_decay=0) # schedulers - scheduler_gen = getattr(torch.optim.lr_scheduler, c.lr_scheduler_gen) - scheduler_disc = getattr(torch.optim.lr_scheduler, c.lr_scheduler_disc) - scheduler_gen = scheduler_gen(optimizer_gen, **c.lr_scheduler_gen_params) - scheduler_disc = scheduler_disc(optimizer_disc, **c.lr_scheduler_disc_params) + scheduler_gen = None + scheduler_disc = None + if 'lr_scheduler_gen' in c: + scheduler_gen = getattr(torch.optim.lr_scheduler, c.lr_scheduler_gen) + scheduler_gen = scheduler_gen(optimizer_gen, **c.lr_scheduler_gen_params) + if 'lr_scheduler_disc' in c: + scheduler_disc = getattr(torch.optim.lr_scheduler, c.lr_scheduler_disc) + scheduler_disc = scheduler_disc(optimizer_disc, **c.lr_scheduler_disc_params) # setup criterion criterion_gen = GeneratorLoss(c)