class CustomClass(PreTrainedModel):

    class_config = PretrainedConfig   # class attribute like this here?

    def __init__(self, config, num_labels):
        super().__init__(config, num_labels)
        self.distilbert = DistilBertModel.from_pretrained('distilbert-base-uncased')
        self.pre_classifier = torch.nn.Linear(768, 768)
        self.dropout = torch.nn.Dropout(0.1)
        self.classifier = torch.nn.Linear(768, num_labels)

    def forward(self, input_ids, attention_mask):
        distilbert_output = self.distilbert(input_ids=input_ids, attention_mask=attention_mask)
        hidden_state = distilbert_output[0]
        pooled_output = hidden_state[:, 0]
        pooled_output = self.pre_classifier(pooled_output)
        pooled_output = torch.nn.Tanh()(pooled_output) 
        pooled_output = self.dropout(pooled_output)  
        output = self.classifier(pooled_output)  
        return output

I added it like above but still no luck to load it with
new_config=PretrainedConfig.from_pretrained('./PATH/') 

new_model=CusomClass.from_pretrained('./PATH/', config=new_config)

Any notebook that I can look at for reference. This is a very common pattern I suppose.