@LearnToGrow are still struggling? It’s working for me. I can save my my custom model with my custom head and are able to load it afterwards with the correct weights.
Important is that you use keyword arguments. So you need to write:
new_model=CustomClass.from_pretrained(pretrained_model_name_or_path='./PATH/', config=new_config)
instead of only:
new_model=CustomClass.from_pretrained('./PATH/', new_config)
And I use a custom config which I assume will be necessary for all custom models. https://huggingface.co/docs/transformers/custom_models
Hope that helps!