File tree Expand file tree Collapse file tree 1 file changed +6
-2
lines changed
mllib/src/main/scala/org/apache/spark/mllib/tree Expand file tree Collapse file tree 1 file changed +6
-2
lines changed Original file line number Diff line number Diff line change @@ -177,9 +177,10 @@ object GradientBoostedTrees extends Logging {
177177 treeStrategy.assertValid()
178178
179179 // Cache input
180- if (input.getStorageLevel == StorageLevel .NONE ) {
180+ val persistedInput = if (input.getStorageLevel == StorageLevel .NONE ) {
181181 input.persist(StorageLevel .MEMORY_AND_DISK )
182- }
182+ true
183+ } else false
183184
184185 timer.stop(" init" )
185186
@@ -265,6 +266,9 @@ object GradientBoostedTrees extends Logging {
265266
266267 logInfo(" Internal timing for DecisionTree:" )
267268 logInfo(s " $timer" )
269+
270+ if (persistedInput) input.unpersist()
271+
268272 if (validate) {
269273 new GradientBoostedTreesModel (
270274 boostingStrategy.treeStrategy.algo,
You can’t perform that action at this time.
0 commit comments