diff options
Diffstat (limited to 'mllib/src/test')
| -rw-r--r-- | mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala b/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala index bd5bd17147..b650a9f092 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala @@ -131,7 +131,7 @@ private[ml] object TreeTests extends SparkFunSuite { * Check if the two models are exactly the same. * If the models are not equal, this throws an exception. */ - def checkEqual(a: TreeEnsembleModel, b: TreeEnsembleModel): Unit = { + def checkEqual[M <: DecisionTreeModel](a: TreeEnsembleModel[M], b: TreeEnsembleModel[M]): Unit = { try { a.trees.zip(b.trees).foreach { case (treeA, treeB) => TreeTests.checkEqual(treeA, treeB) |