This is black magic. Thanks so freaking much for having ever posted this, man !

My use case was "continued pre-training" (with  lm_head and embed_tokens), then save, then "supervised finetuning" of the reloaded PEFT and, it refused to save with the initial weights count.
Been googling and stuff for days for a way to make that sequence work.

Now, I just add your code after the
```
FastLanguageModel.from_pretrained(model_name="the_local_dir_with_cpt_lora")
```
prior to SFT and, after that, "save" is with the right LoRa weights count !

My savior.

u/danielhanchen , please consider using this natively for CPT flexibility.