Fix unit test

Signed-off-by: Adam Li <adam2392@gmail.com>
neurodata · Aug 8, 2024 · 70139e7 · 70139e7
1 parent d91f873
commit 70139e7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/treeple/tree/_honest_tree.py b/treeple/tree/_honest_tree.py
@@ -781,9 +781,9 @@ def _set_leaf_nodes(self, leaf_ids, y, sample_weight):
             total_n_node_samples += y_weight
 
             if leaf_id in seen_leaf_ids:
-                self.tree_.value[leaf_id][:, yval] += 1  # y_weight
+                self.tree_.value[leaf_id][:, yval] += y_weight
             else:
-                self.tree_.value[leaf_id][:, yval] = 1  # y_weight
+                self.tree_.value[leaf_id][:, yval] = y_weight
             seen_leaf_ids.add(leaf_id)
 
         for leaf_id in seen_leaf_ids: