change optimizer initialization for compatibility with Hifi-GAN official implementation

This commit is contained in:
Edresson 2021-04-23 07:54:39 -03:00
parent a6cd04456e
commit d2b6326b8b
1 changed files with 6 additions and 1 deletions

View File

@ -5,6 +5,7 @@
import os
import sys
import time
import itertools
import traceback
from inspect import signature
@ -495,7 +496,11 @@ def main(args): # pylint: disable=redefined-outer-name
optimizer_gen = getattr(torch.optim, c.optimizer)
optimizer_gen = optimizer_gen(model_gen.parameters(), lr=c.lr_gen, **c.optimizer_params)
optimizer_disc = getattr(torch.optim, c.optimizer)
optimizer_disc = optimizer_disc(model_disc.parameters(), lr=c.lr_disc, **c.optimizer_params)
if c.discriminator_model == 'hifigan_discriminator':
optimizer_disc = optimizer_disc(itertools.chain(model_disc.msd.parameters(), model_disc.mpd.parameters()), lr=c.lr_disc, **c.optimizer_params)
else:
optimizer_disc = optimizer_disc(model_disc.parameters(), lr=c.lr_disc, **c.optimizer_params)
# schedulers
scheduler_gen = None