From 8b62f92cc963e2d970a8f27678575be1f17cc7f2 Mon Sep 17 00:00:00 2001
From: QinbinLi <liqinbin1998@gmail.com>
Date: Mon, 11 Jul 2022 11:09:39 +0800
Subject: [PATCH] fix issue #40

---
 python/setup.py                |  2 +-
 src/FedTree/scikit_fedtree.cpp | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/python/setup.py b/python/setup.py
index cedd6338..9c9e9706 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -19,7 +19,7 @@
     copyfile(lib_path, path.join(dirname, "fedtree", path.basename(lib_path)))
 
 setuptools.setup(name="fedtree",
-                 version="1.0.3",
+                 version="1.0.4",
                  packages=["fedtree"],
                  package_dir={"python": "fedtree"},
                  description="A federated learning library for trees",
diff --git a/src/FedTree/scikit_fedtree.cpp b/src/FedTree/scikit_fedtree.cpp
index c01f1ff5..5ba1eb84 100644
--- a/src/FedTree/scikit_fedtree.cpp
+++ b/src/FedTree/scikit_fedtree.cpp
@@ -234,6 +234,11 @@ extern "C" {
                 test_dataset.label.emplace_back(group_label[i]);
             }
         }
+        else{
+            for (int i = 0; i < num_class; ++i) {
+                test_dataset.label.emplace_back(i);
+            }
+        }
         // predict
         SyncArray<float_type> y_predict;
         vector<vector<Tree>> boosted_model_in_mem;
@@ -274,8 +279,16 @@ extern "C" {
         test_dataset.load_from_sparse(row_size, val, row_ptr, col_ptr, NULL, group, num_group, model_param);
         set_logger(verbose);
         test_dataset.label.clear();
-        for (int i = 0; i < num_class; ++i) {
-            test_dataset.label.emplace_back(group_label[i]);
+        if(group_label != NULL) {
+            test_dataset.label.clear();
+            for (int i = 0; i < num_class; ++i) {
+                test_dataset.label.emplace_back(group_label[i]);
+            }
+        }
+        else{
+            for (int i = 0; i < num_class; ++i) {
+                test_dataset.label.emplace_back(i);
+            }
         }
         // predict
         SyncArray<float_type> y_predict;