fix a bug: start training after set weight (#550)

alibaba · Mar 24, 2023 · d2e7d08 · d2e7d08
1 parent 09edda6
commit d2e7d08
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/federatedscope/vertical_fl/xgb_base/worker/train_wrapper.py b/federatedscope/vertical_fl/xgb_base/worker/train_wrapper.py
@@ -178,6 +178,7 @@ def callback_funcs_for_local_best_gain(self, message: Message):
                     self.comm_manager.send(send_message)
             else:
                 self.trainer._set_weight_and_status(tree_num, node_num)
+                self.train(tree_num=tree_num, node_num=node_num + 1)
 
     # Bind method to instance
     client.train = types.MethodType(train, client)