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 authored and GitHub committed 2021-01-21 21:16:30 +01:00
1 parent b70bef579a
commit f251dc8c0e
1 file changed
+1 -1
+1 -1
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)