Update flwr client context usage

NVIDIA · Oct 7, 2024 · 9405342 · 9405342
1 parent fa1fd2b
commit 9405342
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/examples/hello-world/hello-flower/flwr-pt-tb/flwr_pt_tb/client.py b/examples/hello-world/hello-flower/flwr-pt-tb/flwr_pt_tb/client.py
@@ -36,18 +36,16 @@ class FlowerClient(NumPyClient):
     def __init__(self, context: Context):
         super().__init__()
         self.writer = SummaryWriter()
-        self.set_context(context)
+        self.flwr_context = context
+
         if "step" not in context.state.metrics_records:
             self.set_step(0)
 
     def set_step(self, step: int):
-        context = self.get_context()
-        context.state = RecordSet(metrics_records={"step": MetricsRecord({"step": step})})
-        self.set_context(context)
+        self.flwr_context.state = RecordSet(metrics_records={"step": MetricsRecord({"step": step})})
 
     def get_step(self):
-        context = self.get_context()
-        return int(context.state.metrics_records["step"]["step"])
+        return int(self.flwr_context.state.metrics_records["step"]["step"])
 
     def fit(self, parameters, config):
         step = self.get_step()