Update model.py

model.py

@@ -1,23 +1,31 @@
 import torch
 import torch.nn as nn
-from torchvision.models import swin_t, Swin_T_Weights
-from huggingface_hub import PyTorchModelHubMixin
+from torchvision.models import swin_t
+from transformers import PretrainedConfig, PreTrainedModel
 
-class SwinClassifier(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, num_classes):
-        super().__init__()
-        # 1. Load the backbone
+# 1. Define a Config class
+class SwinClassifierConfig(PretrainedConfig):
+    model_type = "swin_classifier"
+    def __init__(self, num_classes=18, **kwargs):
+        super().__init__(**kwargs)
+        self.num_classes = num_classes
+
+# 2. Update your Model class to inherit from PreTrainedModel
+class SwinClassifier(PreTrainedModel):
+    config_class = SwinClassifierConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        # Use config.num_classes instead of a raw number
         self.backbone = swin_t()
-        
-        # 2. Get features and replace the head
         num_features = self.backbone.head.in_features
         
-        # We replace the original head with our custom Sequential block
         self.backbone.head = nn.Sequential(
             nn.Linear(num_features, 256),
             nn.ReLU(inplace=True),
             nn.Dropout(0.5),
-            nn.Linear(256, num_classes)
+            # Use the value from the config
+            nn.Linear(256, config.num_classes)
         )
 
     def forward(self, x):