Getting size mismatch when loading checkpoints saved py the code - tested with PHI-2

[sowmaster]

Hi,
Thank you for the great work! I am having issues when i try to load the checkpoints saved during alignment for evaluation.
I loaded the checkpoints with the code (just regular loading which is also used in alpacaEval with HF completions):

model = AutoModelForCausalLM.from_pretrained(path_to_ckpt, cache_dir=cache_dir, **model_kwargs).eval()

The following error show up which says there is mismatch in in param dimensions.

RuntimeError: Error(s) in loading state_dict for PhiForCausalLM:
        size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([65550080]) from checkpoint, the shape in current model is torch.S
ize([51200, 2560]).
        size mismatch for model.final_layernorm.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size(
[2560]).
        size mismatch for model.final_layernorm.bias: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2
560]).
        size mismatch for lm_head.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([51200, 2560]).

Could be related to how the ckpt are saved. Any ideas of a fix?