adds default best_iteration in more spots

ClimbsRocks · Feb 13, 2018 · 1aedcdb · 1aedcdb
1 parent d37b0fe
commit 1aedcdb
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/auto_ml/utils_model_training.py b/auto_ml/utils_model_training.py
@@ -621,13 +621,11 @@ def predict(self, X, verbose=False):
 
 
         if self.model_name[:4] == 'LGBM':
-            best_iteration = 0
+            best_iteration = -1
             try:
                 best_iteration = self.model.best_iteration_
             except AttributeError:
                 best_iteration = self.model.best_iteration
-            if best_iteration is None:
-                best_iteration = 0
             predictions = self.model.predict(X, num_iteration=best_iteration)
         else:
             predictions = self.model.predict(X_predict)