cast to np.array before prediction (#20)

intro-stat-learning · Jun 4, 2024 · a956573 · a956573
1 parent 7bb8fa4
commit a956573
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ISLP/bart/bart.py b/ISLP/bart/bart.py
@@ -105,7 +105,7 @@ def predict(self,
 
         for trees in self.trees_sample_:
             for tree in trees:
-                tree_fit = np.array([tree.predict_out_of_sample(x) for x in X])
+                tree_fit = np.array([tree.predict_out_of_sample(x) for x in np.asarray(X)])
                 output += tree_fit
         output = output / nsample
         return self._inverse(output)
@@ -122,7 +122,7 @@ def staged_predict(self,
 
         for nstep, trees in enumerate(trees_sample_):
             for tree in trees:
-                tree_fit = np.array([tree.predict_out_of_sample(x) for x in X])
+                tree_fit = np.array([tree.predict_out_of_sample(x) for x in np.asarray(X)])
                 output[nstep] += tree_fit
 
         output = np.cumsum(output, 0) / (np.arange(nsample) + 1)[:,None]