Update app.py

app.py

@@ -147,13 +147,20 @@ class SkinViT(nn.Module):
 class DermNetViT(nn.Module):
     def __init__(self, num_classes):
         super(DermNetViT, self).__init__()
+        from torchvision.models import vit_l_16, ViT_L_16_Weights
         self.model = vit_l_16(weights=ViT_L_16_Weights.DEFAULT)
-        in_features = self.model.heads.head.in_features
-        self.model.heads.head = nn.Linear(in_features, num_classes)
+        in_features = self.model.heads[0].in_features
+        # Match the structure used during training (with head.1)
+        self.model.heads = nn.Sequential(
+            nn.Linear(in_features, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, num_classes)
+        )
 
     def forward(self, x):
         return self.model(x)
 
+
 #multilabel_model = torch.load("D:/DR/RAG/BestModels2703/skin_vit_fold10.pth", map_location='cpu')
 #multiclass_model = torch.load("D:/DR/RAG/BestModels2703/best_dermnet_vit.pth", map_location='cpu')
 
@@ -163,8 +170,10 @@ multiclass_model_path = hf_hub_download(repo_id="santhoshraghu/DermBOT", filenam
 
 multilabel_model = SkinViT(num_classes=len(multilabel_class_names))
 multiclass_model = DermNetViT(num_classes=len(multiclass_class_names))
+
 multilabel_model.load_state_dict(torch.load(multilabel_model_path, map_location="cpu"))
 multiclass_model.load_state_dict(torch.load(multiclass_model_path, map_location="cpu"))
+
 multilabel_model.eval()
 multiclass_model.eval()