models/student_model.py

from transformers import DistilBertForSequenceClassification

def get_student_model(config):
    model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
    model.config.hidden_size = config['model']['student']['hidden_size']
    return model