Update app.py

app.py

@@ -39,22 +39,13 @@ def init_db():
 
 init_db()
 
-model = models.densenet169(pretrained=False)
-model.classifier = nn.Linear(model.classifier.in_features, 3)
+densenet_model = models.densenet169(pretrained=False)
+densenet_model.classifier = nn.Linear(densenet_model.classifier.in_features, 3)
 
-# 2️⃣ Load your checkpoint
+model = torch.nn.Sequential(densenet_model)  # ✅ IMPORTANT
 state_dict = torch.load('densenet169_seed40_best.pt', map_location='cpu')
+model.load_state_dict(state_dict)
 
-# 3️⃣ Strip the "features.0." prefix
-new_state_dict = {}
-for k, v in state_dict.items():
-    new_key = k.replace("features.0.", "")
-    new_state_dict[new_key] = v
-
-# 4️⃣ Load the stripped state dict
-model.load_state_dict(new_state_dict)
-
-# 5️⃣ Set eval mode
 model.eval()
 
 print("✅ Model successfully loaded and ready for inference.")