From d65def13c22736bd16cfdcfe167c936c5ad9353f Mon Sep 17 00:00:00 2001 From: kta-intel Date: Mon, 18 Nov 2024 12:29:07 -0800 Subject: [PATCH] raise value error when model is empty when trying to set tensor dict Signed-off-by: kta-intel --- openfl/federated/task/runner_xgb.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/openfl/federated/task/runner_xgb.py b/openfl/federated/task/runner_xgb.py index f65b1039fc..5169335614 100644 --- a/openfl/federated/task/runner_xgb.py +++ b/openfl/federated/task/runner_xgb.py @@ -315,9 +315,14 @@ def set_tensor_dict(self, tensor_dict, with_opt_vars=False): """ # The with_opt_vars argument is not used in this method self.global_model = tensor_dict["local_tree"] - global_model_byte_array = bytearray(self.global_model.astype(np.uint8).tobytes()) - self.bst = xgb.Booster() - self.bst.load_model(global_model_byte_array) + if ( + isinstance(self.global_model, np.ndarray) and self.global_model.size == 0 + ) or self.global_model is None: + raise ValueError("The model does not exist or is empty.") + else: + global_model_byte_array = bytearray(self.global_model.astype(np.uint8).tobytes()) + self.bst = xgb.Booster() + self.bst.load_model(global_model_byte_array) def save_native( self,