From 6a1f2d242c767241205edb16a1bc7dd295325c49 Mon Sep 17 00:00:00 2001 From: James Lamb Date: Mon, 25 Jan 2021 09:45:45 -0600 Subject: [PATCH] fix get_combined_model --- python/cuml/dask/ensemble/base.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/python/cuml/dask/ensemble/base.py b/python/cuml/dask/ensemble/base.py index 3a5e0e898a..af8eca8909 100644 --- a/python/cuml/dask/ensemble/base.py +++ b/python/cuml/dask/ensemble/base.py @@ -272,7 +272,18 @@ def get_combined_model(self): if self._get_internal_model() is None: self._set_internal_model(self._concat_treelite_models()) - return BaseEstimator.get_combined_model(self) + internal_model = self._check_internal_model(self._get_internal_model()) + + if isinstance(self.internal_model, Iterable): + # This function needs to return a single instance of cuml.Base, + # even if the class is just a composite. + raise ValueError("Expected a single instance of cuml.Base " + "but got %s instead." % type(self.internal_model)) + + elif isinstance(self.internal_model, Future): + internal_model = self.internal_model.result() + + return internal_model def _func_fit(model, input_data, convert_dtype):