Skip to content

Save new `model` when fine-tuning

When --train.pretrain true PyLaia loads the pretrained model, and the last layer is adapted to the new charset.

After training, the best checkpoint is saved, but the new model is not saved.

When decoding, PyLaia throws a size mismatch error because the checkpoint is not compatible with the model.

size mismatch for linear.weight: copying a param with shape torch.Size([169, 512]) from checkpoint, the shape in current model is torch.Size([91, 512]).\n\tsize mismatch for linear.bias: copying a param with shape torch.Size([169]) from checkpoint, the shape in current model is torch.Size([91]).'

To resolve this issue, we need to save the new model object after the last layer is reinitialized.