Oh you need to add some class variables to your custom model, specifically:
config_class = CustomConfig

and have your config be an instance of a CustomConfig (it might work with config_class = PretrainedConfig but I’m not 100% sure).