my bad

model.py

@@ -11,14 +11,16 @@ def create_model():
     weights = models.Swin_B_Weights.DEFAULT
     transform = weights.transforms()
 
-    model = models.efficientnet_b2(weights = weights)
+    model = models.swin_b(weights = weights)
 
     for param in model.parameters():
         param.requires_grad = False 
 
-    model.classifier = nn.Sequential(
-        nn.Dropout(p = 0.3, inplace = True),
-        nn.Linear(in_features = 1408, out_features = 101, bias = True)
+    model.head = nn.Sequential(
+        # nn.Linear(in_features = 1024, out_features = 512, bias = True),
+        # nn.ReLU(),
+        # nn.Dropout(p = 0.3, inplace = True),
+        nn.Linear(in_features = 1024, out_features = num_classes, bias = True)
     )