Parameter name inconsistency between 1st and 2nd stage

[halsay]

StyleTTS2/train_second.py

Line 39 in 2c427fc

return getattr(self.module, name)

In 2nd stage training, names of all parameters are added 'module.' as a prefix. However, in the 1st stage they are not.
In that case, parameters trained from 1st stage are not really loaded in 2nd stage and you won't get an error because they are not strict loaded (It even prints the model is loaded).
https://github.com/yl4579/StyleTTS2/blob/main/models.py#L696
I get Nan at the beginning of 2nd stage training and I manage to find out the problem described above. And 2nd stage training becomes normal after the parameters successfully loaded.

[RillmentGames]

I ran into this as well, Issue #21 has the solution for going from stage 1 to stage 2 but you need to remove the extra code again when reloading a stage2 checkpoint.

[yl4579]

You will need to run the first stage with multiple GPUs. Currently single GPU training for first stage is not supported. You can also remove all MyDataParallel in train_second.py to fix this problem if you only have one GPU.