zeroed tensor at end of Flatten

interpretml · Dec 25, 2024 · 71ca3a7 · 71ca3a7
1 parent 748cc6b
commit 71ca3a7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/shared/libebm/PartitionOneDimensionalBoosting.cpp b/shared/libebm/PartitionOneDimensionalBoosting.cpp
@@ -350,8 +350,10 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
 
                EBM_ASSERT(bNominal || pSplit == cSlices - 1 + pInnerTermUpdate->GetSplitPointer(iDimension));
 
+               pInnerTermUpdate->SetCountSlices(iDimension, 1);
+
                LOG_0(Trace_Verbose, "Exited Flatten");
-               return Error_IllegalParamVal;
+               return Error_None;
             }
             if(!pTreeNode->DECONSTRUCT_IsRightChildTraversal()) {
                // we checked earlier that countBins could be converted to a UIntSplit