handls inconsistency in lgb api- forces best_iteration to be -1 inste…

…ad of None, even if they give us None
ClimbsRocks · Feb 13, 2018 · 54da719 · 54da719
1 parent 1aedcdb
commit 54da719
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/auto_ml/utils_model_training.py b/auto_ml/utils_model_training.py
@@ -626,6 +626,10 @@ def predict(self, X, verbose=False):
                 best_iteration = self.model.best_iteration_
             except AttributeError:
                 best_iteration = self.model.best_iteration
+            # At version 2.1.0, lgbm returns best_iteration = None, but doesn't take None as an input for num_iteration. this issue only exists in python3(.6)
+            if best_iteration is None:
+                best_iteration = -1
+
             predictions = self.model.predict(X, num_iteration=best_iteration)
         else:
             predictions = self.model.predict(X_predict)