Try deberta small

app.py

@@ -46,7 +46,7 @@ student_model_config = teacher_model.config
 student_model_config.num_attention_heads = 6 
 student_model_config.num_hidden_layers = 4
 student_model = DebertaV2ForTokenClassification.from_pretrained(
-    "microsoft/deberta-v3-base",
+    "microsoft/deberta-v3-small",
     config=student_model_config, ignore_mismatched_sizes=True)
 print(student_model)
 print_trainable_parameters(student_model)
@@ -58,7 +58,7 @@ if torch.cuda.is_available():
 # Load data. 
 raw_dataset = load_dataset("ai4privacy/pii-masking-400k", split='train')
 raw_dataset = raw_dataset.filter(lambda example: example["language"].startswith("en"))
-raw_dataset = raw_dataset.select(range(2000, 4000)) 
+raw_dataset = raw_dataset.select(range(4000)) 
 raw_dataset = raw_dataset.train_test_split(test_size=0.2)
 print(raw_dataset)
 print(raw_dataset.column_names)