Update app.py

app.py

@@ -39,13 +39,23 @@ def init_db():
 
 init_db()
 
-densenet_model = models.densenet169(pretrained=False)
-densenet_model.classifier = nn.Linear(densenet_model.classifier.in_features, 3)
+model = models.densenet169(pretrained=False)
+model.classifier = nn.Linear(model.classifier.in_features, 3)
 
-model = torch.nn.Sequential(densenet_model)  # ✅ IMPORTANT
+# 2️⃣ Load the checkpoint
 state_dict = torch.load('densenet169_seed40_best.pt', map_location='cpu')
-model.load_state_dict(state_dict)
 
+# 3️⃣ Rename keys
+new_state_dict = {}
+for k, v in state_dict.items():
+    if k.startswith('0.'):
+        new_key = k[2:]  # Removes the leading '0.'
+    else:
+        new_key = k
+    new_state_dict[new_key] = v
+
+# 4️⃣ Load
+model.load_state_dict(new_state_dict)
 model.eval()
 
 print("✅ Model successfully loaded and ready for inference.")