Support Python 3.7 in iterate_leaf_indexes and staged_predict fix
@@ -1812,7 +1812,10 @@ class CatBoost(_CatBoostBase):
ntree_end = self.tree_count_
staged_predict_iterator = self._staged_predict_iterator(data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose)
while True:
+ try:
predictions = staged_predict_iterator.next()
+ except StopIteration:
+ return
yield predictions[0] if data_is_single_object else predictions
def staged_predict(self, data, prediction_type='RawFormulaVal', ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None):
@@ -1872,7 +1875,10 @@ class CatBoost(_CatBoostBase):
data, _ = self._process_predict_input_data(data, "iterate_leaf_indexes")
leaf_indexes_iterator = self._leaf_indexes_iterator(data, ntree_start, ntree_end)
while True:
+ try:
yield leaf_indexes_iterator.next()
+ except StopIteration:
+ return
def iterate_leaf_indexes(self, data, ntree_start=0, ntree_end=0):
"""
