Key Error "Iteration" in checkpoint path

[shehrum]

I am trying to train on a custom dataset following the steps for

"Training using a pre-trained model"

with the command:

python train.py -c config.json -p train_config.ignore_layers=["speaker_embedding.weight"] train_config.checkpoint_path="models/flowtron_ljs.pt"

However, I get the error:

File "train.py", line 110, in load_checkpoint iteration = checkpoint_dict['iteration'] KeyError: 'iteration'

It seems as the model doesn't have these keys. How should I go about fixing/modifying the training script to cater for that. Since the current train.py script throws this error. Would appreciate if anyone could help fix the issue.

[stqc]

in the train script edit the code to this

try: iteration = checkpoint_dict['iteration'] model_dict = checkpoint_dict['model'].state_dict() except: iteration = 1 model_dict = checkpoint_dict['state_dict']

[shehrum]

Thankyou! Changed to this already and it worked.