# Load student model
my_config = DistilBertConfig(n_heads=8, n_layers=4) # drop 4 heads per layer and 2 layers

student_model = DistilBertForSequenceClassification
                                    .from_pretrained("distilbert-base-uncased",
                                    config=my_config,)
                                    .to(device)