Update train_tacotron.py

When attempting to fine-tune a model with "prenet_type": "bn" that was originally trained with "prenet_type": "original", a RuntimeError is thrown that stops the training.

By catching the RuntimeError, the required layers can be partially restored and the training will continue without any problems.
This commit is contained in:
Alexander Korolev 2021-01-21 21:16:30 +01:00 committed by GitHub
parent b70bef579a
commit f251dc8c0e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -549,7 +549,7 @@ def main(args): # pylint: disable=redefined-outer-name
scaler.load_state_dict(checkpoint["scaler"])
if c.reinit_layers:
raise RuntimeError
except KeyError:
except (KeyError, RuntimeError):
print(" > Partial model initialization.")
model_dict = model.state_dict()
model_dict = set_init_dict(model_dict, checkpoint['model'], c)