Trying to accomodate new API v2

benchmark_utils/template_flamby_strategy.py

@@ -93,8 +93,10 @@ def run(self, callback):
         # We are reproducing the run method but this time a callback checks
         # stopping-criterion at each round, which allows to cache computations
         # and do a single run
-        while callback(strat.models_list[0].model):
+        self.final_model = strat.models_list[0].model
+        while callback():
             strat.perform_round()
+            self.final_model = strat.models_list[0].model
 
         self.final_model = strat.models_list[0].model
 
@@ -103,7 +105,7 @@ def get_result(self):
         # The outputs of this function are the arguments of `Objective.compute`
         # This defines the benchmark's API for solvers' results.
         # it is customizable for each benchmark.
-        return self.final_model
+        return {"model": self.final_model}
 
     # Not used if callback is used
     @staticmethod

objective.py

@@ -256,7 +256,7 @@ def robust_metric(y_true, y_pred):
     def get_one_result(self):
         # Return one solution. The return value should be an object compatible
         # with `self.compute`. This is mainly for testing purposes.
-        return self.model_arch()
+        return dict(model=self.model_arch())
 
     def get_objective(self):
         # Define the information to pass to each solver to run the benchmark.