diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py index a96a8f0299a9..66538626ced3 100644 --- a/python-package/lightgbm/basic.py +++ b/python-package/lightgbm/basic.py @@ -1267,6 +1267,10 @@ def __setstate__(self, state): state['handle'] = handle self.__dict__.update(state) + def free_dataset(self): + self.__dict__.pop('train_set', None) + self.__dict__.pop('valid_sets', None) + def set_train_data_name(self, name): self.__train_data_name = name diff --git a/python-package/lightgbm/sklearn.py b/python-package/lightgbm/sklearn.py index 1101ede481d6..26a9797eff04 100644 --- a/python-package/lightgbm/sklearn.py +++ b/python-package/lightgbm/sklearn.py @@ -418,6 +418,10 @@ def get_meta_data(collection, i): if early_stopping_rounds is not None: self.best_iteration = self._Booster.best_iteration self.best_score = self._Booster.best_score + + # free dataset + self.booster_.free_dataset() + del train_set, valid_sets return self def predict(self, X, raw_score=False, num_iteration=0):