update app.py 2023-08-31 07:48

app.py

@@ -40,8 +40,13 @@ 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)
-
+densenet.load_state_dict(model_weights,strict=False)
+'''
+weights = {k: torch.from_numpy(v).to(self.device) if isinstance(v, np.ndarray) else v.to(self.device) for k, v in weights.items()}
+# creat new state_dict and del fc.weight andfc.bias
+new_state_dict = {k: v for k, v in weights.items() if k not in ["fc.weight", "fc.bias"]}
+self.model.load_state_dict(new_state_dict, strict=False)
+'''
 
 def predict(img) -> Tuple[Dict, float]:
     """Transforms and performs a prediction on img and returns prediction and time taken.