update app.py

app.py

@@ -39,7 +39,8 @@ densenet, densenet_transforms = create_densenet121_model(num_classes=1)
 
 # Load saved weights
 densenet.load_state_dict(torch.load("FL_global_model.pt", map_location=torch.device("cpu")))
-
+model_weights = state_dict["model"]
+densenet.load_state_dict(model_weights)
 
 
 def predict(img) -> Tuple[Dict, float]:
@@ -49,13 +50,13 @@ def predict(img) -> Tuple[Dict, float]:
     start_time = timer()
 
     # Transform the target image and add a batch dimension
-    img = effnetb2_transforms(img).unsqueeze(0)
+    img = densenet_transforms(img).unsqueeze(0)
 
     # Put model into evaluation mode and turn on inference mode
-    effnetb2.eval()
+    densenet.eval()
     with torch.inference_mode():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
-        pred_probs = torch.sigmoid(effnetb2(img)).squeeze()
+        pred_probs = torch.sigmoid(densenet(img)).squeeze()
 
     # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
     pred_labels_and_probs = {