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