fixes bug with memory_optimized lgbm

ClimbsRocks · Feb 22, 2018 · 0519678 · 0519678
1 parent 9df56a8
commit 0519678
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/auto_ml/utils_model_training.py b/auto_ml/utils_model_training.py
@@ -476,11 +476,15 @@ def predict_proba(self, X, verbose=False):
                     best_iteration = self.model.best_iteration
                 except AttributeError:
                     best_iteration = self.model.best_iteration_
-                predictions = self.model.predict_proba(X, num_iteration=best_iteration)
+                try:
+                    predictions = self.model.predict_proba(X, num_iteration=best_iteration)
+                except AttributeError as e:
+                    predictions = self.model.predict(X, num_iteration=best_iteration)
             else:
                 predictions = self.model.predict_proba(X)
 
         except AttributeError as e:
+
             try:
                 predictions = self.model.predict(X)
             except TypeError as e:
@@ -501,8 +505,10 @@ def predict_proba(self, X, verbose=False):
             for prediction in predictions:
                 if prediction == 1:
                     tupled_predictions.append([0,1])
-                else:
+                elif prediction == 0:
                     tupled_predictions.append([1,0])
+                else:
+                    tupled_predictions.append([1 - prediction, prediction])
             predictions = tupled_predictions