Fix _VotingBooster

Y-oHr-N · Feb 20, 2020 · f3996fe · f3996fe
1 parent 685d8d9
commit f3996fe
Showing 1 changed file with 6 additions and 16 deletions.
diff --git a/optgbm/sklearn.py b/optgbm/sklearn.py
@@ -249,25 +249,15 @@ def from_representations(
 
         return cls(boosters, weights=weights)
 
-    def predict(
-        self, X: TWO_DIM_ARRAYLIKE_TYPE, **kwargs: Any
-    ) -> TWO_DIM_ARRAYLIKE_TYPE:
-        results = []
-
-        for b in self.boosters:
-            result = b.predict(X, **kwargs)
-
-            results.append(result)
-
-        return np.average(results, axis=0, weights=self.weights)
-
     def feature_importance(self, **kwargs: Any) -> np.ndarray:
-        results = []
+        results = [b.feature_importance(**kwargs) for b in self.boosters]
 
-        for b in self.boosters:
-            result = b.feature_importance(**kwargs)
+        return np.average(results, axis=0, weights=self.weights)
 
-            results.append(result)
+    def predict(
+        self, X: TWO_DIM_ARRAYLIKE_TYPE, **kwargs: Any
+    ) -> TWO_DIM_ARRAYLIKE_TYPE:
+        results = [b.predict(X, **kwargs) for b in self.boosters]
 
         return np.average(results, axis=0, weights=self.weights)