update models.py

models.py

@@ -2,6 +2,17 @@ import torch
 from torch import nn
 from transformers import PreTrainedModel
 
+from transformers import PretrainedConfig
+
+class CustomClassificationConfig(PretrainedConfig):
+    model_type = "custom_classifier"
+
+    def __init__(self, input_dim=32, hidden_dim=64, num_classes=2, **kwargs):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_classes = num_classes
+        
 class CustomClassifier(PreTrainedModel):
     config_class = CustomClassificationConfig