Merge branch 'torch-update' of https://github.com/tanwarsh/openfl int…

…o torch-update
securefederatedai · Jan 21, 2025 · f959dc3 · f959dc3
2 parents 932151c + ea20dde
commit f959dc3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/openfl/component/aggregator/aggregator.py b/openfl/component/aggregator/aggregator.py
@@ -352,6 +352,7 @@ def _save_model(self, round_number, file_path):
         ]
         tensor_dict = {}
         tensor_tuple_dict = {}
+        next_round_tensors = {}
         for tk in tensor_keys:
             tk_name, _, _, _, _ = tk
             tensor_value = self.tensor_db.get_tensor_from_cache(tk)
@@ -373,7 +374,7 @@ def _save_model(self, round_number, file_path):
                         self.next_model_round_number, ("model",)
                     )
                 self.persistent_db.finalize_round(
-                    tensor_tuple_dict, next_round_tensors, self.round_number, self.best_model_score
+                    tensor_tuple_dict, next_round_tensors, round_number, self.best_model_score
                 )
                 logger.info(
                     "Persist model and clean task result for round %s",