Update app.py

app.py

@@ -40,16 +40,19 @@ def init_db():
 init_db()
 
 model = models.densenet169(pretrained=False)
-
-# IMPORTANT: Adjust final layer to match your number of classes (3 in your case)
 model.classifier = nn.Linear(model.classifier.in_features, 3)
 
-# Load state dict
-state_dict = torch.load('densenet169_seed40_best.pt', map_location=torch.device('cpu'))
-model.load_state_dict(state_dict)
+# 2️⃣ Load checkpoint
+checkpoint = torch.load('densenet169_seed40_best.pt', map_location='cpu')
+
+# ✅ If checkpoint contains state_dict directly
+# model.load_state_dict(checkpoint)
 
-# Set eval mode
-model.eval()
+# ✅ If checkpoint contains a dict
+if 'model_state_dict' in checkpoint:
+    model.load_state_dict(checkpoint['model_state_dict'])
+else:
+    model.load_state_dict(checkpoint)
 
 # ✅ Preprocess Image
 transform = transforms.Compose([