Split stats now list object

JoKoum · Dec 3, 2021 · 662fba7 · 662fba7
1 parent d48a248
commit 662fba7
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/utils/StreamingGradientTree.py b/utils/StreamingGradientTree.py
@@ -108,11 +108,7 @@ def reset(self):
         self.mUpdateStats = GradHessStats()
         self.mInstances = 0
 
-        for i in range(len(self.tree.mFeatureInfo)):
-            self.mSplitStats[i] = [GradHessStats() for _ in range(self.tree.mFeatureInfo[i].categories)]
-
-            for j in range(len(self.mSplitStats[i])):
-                self.mSplitStats[i][j] = GradHessStats()
+        self.mSplitStats = [[GradHessStats() for _ in range(self.tree.mFeatureInfo[i].categories)] for i in range(len(self.tree.mFeatureInfo))]
 
     def getLeaf(self, features):
         if not self.mChildren:
@@ -269,7 +265,7 @@ def applySplit(self, split):
             print("Unhandled attribute type")
 
         # Free up memory used by the split stats
-        self.mSplitStats = {}
+        self.mSplitStats = []
 
     def computeDeltaPrediction(self, gradHess):
         return -gradHess.gradient / (gradHess.hessian + self.tree.mOptions.mLambda + 2.225E-308)