diff --git a/learnware/learnware/reuse.py b/learnware/learnware/reuse.py
index f17515e..d140470 100644
--- a/learnware/learnware/reuse.py
+++ b/learnware/learnware/reuse.py
@@ -643,6 +643,21 @@ class EnsemblePruningReuser(BaseReuser):
 
         return res["Vars"][bst_pop]
 
+    def _get_predict(self, X, selected_idxes):
+        preds = []
+        for idx in selected_idxes:
+            pred_y = self.learnware_list[idx].predict(X)
+            if len(pred_y.shape) == 1:
+                pred_y = pred_y.reshape(-1, 1)
+            elif len(pred_y.shape) == 2:
+                if pred_y.shape[1] > 1:
+                    pred_y = pred_y.argmax(axis=1).reshape(-1, 1)
+            else:
+                raise ValueError("Model output must be a 1D or 2D vector")
+            preds.append(pred_y)
+        
+        return np.concatenate(preds, axis=1)
+    
     def fit(self, val_X: np.ndarray, val_y: np.ndarray, maxgen: int = 500):
         """Ensemble pruning based on the validation set
 
@@ -656,11 +671,7 @@ class EnsemblePruningReuser(BaseReuser):
             The maximum number of iteration rounds in ensemble pruning algorithms.
         """
         # Get the prediction of each learnware on the validation set
-        v_predict = []
-        for idx in range(len(self.learnware_list)):
-            pred_y = self.learnware_list[idx].predict(val_X).reshape(-1, 1)
-            v_predict.append(pred_y)
-        v_predict = np.concatenate(v_predict, axis=1)
+        v_predict = self._get_predict(val_X, list(range(len(self.learnware_list))))
         v_true = val_y.reshape(-1, 1)
 
         # Run ensemble pruning algorithm
@@ -686,13 +697,9 @@ class EnsemblePruningReuser(BaseReuser):
         np.ndarray
             Prediction given by ensemble method
         """
-        preds = []
-        for idx in self.selected_idxes:
-            pred_y = self.learnware_list[idx].predict(user_data).reshape(-1, 1)
-            preds.append(pred_y)
+        preds = self._get_predict(user_data, self.selected_idxes)
 
         if self.mode == "regression":
-            return np.concatenate(preds, axis=1).mean(axis=1)
-        elif self.option == "binary" or self.option == "multiclass":
-            preds = np.concatenate(preds, axis=1)
+            return preds.mean(axis=1)
+        elif self.mode == "binary" or self.mode == "multiclass":
             return np.apply_along_axis(lambda x: np.bincount(x).argmax(), axis=1, arr=preds)
diff --git a/tests/test_workflow/test_workflow.py b/tests/test_workflow/test_workflow.py
index bd3a64b..73417d3 100644
--- a/tests/test_workflow/test_workflow.py
+++ b/tests/test_workflow/test_workflow.py
@@ -12,7 +12,7 @@ from shutil import copyfile, rmtree
 
 import learnware
 from learnware.market import EasyMarket, BaseUserInfo
-from learnware.learnware import JobSelectorReuser, AveragingReuser
+from learnware.learnware import JobSelectorReuser, AveragingReuser, EnsemblePruningReuser
 import learnware.specification as specification
 
 curr_root = os.path.dirname(os.path.abspath(__file__))
@@ -172,15 +172,13 @@ class TestAllWorkflow(unittest.TestCase):
         print("Total Item:", len(easy_market))
 
         X, y = load_digits(return_X_y=True)
-        _, data_X, _, data_y = train_test_split(X, y, test_size=0.3, shuffle=True)
+        train_X, data_X, train_y, data_y = train_test_split(X, y, test_size=0.3, shuffle=True)
 
         stat_spec = specification.utils.generate_rkme_spec(X=data_X, gamma=0.1, cuda_idx=0)
         user_info = BaseUserInfo(semantic_spec=user_semantic, stat_info={"RKMEStatSpecification": stat_spec})
 
         _, _, _, mixture_learnware_list = easy_market.search_learnware(user_info)
 
-        # print("Mixture Learnware:", mixture_learnware_list)
-
         # Based on user information, the learnware market returns a list of learnwares (learnware_list)
         # Use jobselector reuser to reuse the searched learnwares to make prediction
         reuse_job_selector = JobSelectorReuser(learnware_list=mixture_learnware_list)
@@ -189,9 +187,15 @@ class TestAllWorkflow(unittest.TestCase):
         # Use averaging ensemble reuser to reuse the searched learnwares to make prediction
         reuse_ensemble = AveragingReuser(learnware_list=mixture_learnware_list, mode="vote")
         ensemble_predict_y = reuse_ensemble.predict(user_data=data_X)
+        
+        # Use ensemble pruning reuser to reuse the searched learnwares to make prediction
+        reuse_ensemble = EnsemblePruningReuser(learnware_list=mixture_learnware_list, mode="multiclass")
+        reuse_ensemble.fit(train_X[-200:], train_y[-200:])
+        ensemble_pruning_predict_y = reuse_ensemble.predict(user_data=data_X)
 
         print("Job Selector Acc:", np.sum(np.argmax(job_selector_predict_y, axis=1) == data_y) / len(data_y))
-        print("Averaging Selector Acc:", np.sum(np.argmax(ensemble_predict_y, axis=1) == data_y) / len(data_y))
+        print("Averaging Reuser Acc:", np.sum(np.argmax(ensemble_predict_y, axis=1) == data_y) / len(data_y))
+        print("Ensemble Pruning Reuser Acc:", np.sum(ensemble_pruning_predict_y == data_y) / len(data_y))
 
 
 def suite():