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

app.py

@@ -35,10 +35,14 @@ def create_densenet121_model(num_classes: int = 1, seed: int = 42):
     return model, transforms
 
 # Create densenet121 model
-densenet, densenet_transforms = create_densenet121_model(num_classes=1)
+densenet, densenet_transforms = create_densenet121_model(num_classes=2)
 
 # Load saved weights
-densenet.load_state_dict(torch.load("FL_global_model.pt", map_location=torch.device("cpu")))
+# densenet.load_state_dict(torch.load("FL_global_model.pt", map_location=torch.device("cpu")))
+state_dict = torch.load("FL_global_model.pt", map_location=torch.device("cpu"))
+print("==============")
+print(state_dict.keys())
+print("==============")
 model_weights = state_dict["model"]
 densenet.load_state_dict(model_weights,strict=False)
 '''