Update app.py

app.py

@@ -276,6 +276,38 @@ def train(net, trainloader, epochs):
             optimizer.step()
             optimizer.zero_grad()
 
+
+
+
+# class SaveModelStrategy(fl.server.strategy.FedAvg):
+#     def aggregate_fit(
+#         self,
+#         server_round: int,
+#         results: List[Tuple[fl.server.client_proxy.ClientProxy, fl.common.FitRes]],
+#         failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]],
+#     ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:
+#         """Aggregate model weights using weighted average and store checkpoint"""
+
+#         # Call aggregate_fit from base class (FedAvg) to aggregate parameters and metrics
+#         aggregated_parameters, aggregated_metrics = super().aggregate_fit(server_round, results, failures)
+
+#         if aggregated_parameters is not None:
+#             print(f"Saving round {server_round} aggregated_parameters...")
+
+#             # Convert `Parameters` to `List[np.ndarray]`
+#             aggregated_ndarrays: List[np.ndarray] = fl.common.parameters_to_ndarrays(aggregated_parameters)
+
+#             # Convert `List[np.ndarray]` to PyTorch`state_dict`
+#             params_dict = zip(net.state_dict().keys(), aggregated_ndarrays)
+#             state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
+#             net.load_state_dict(state_dict, strict=True)
+
+#             # Save the model
+#             torch.save(net.state_dict(), f"model_round_{server_round}.pth")
+
+#         return aggregated_parameters, aggregated_metrics
+
+        
 def test(net, testloader):
     metric = load_metric("accuracy")
     net.eval()