diff --git a/boolformer/model/model_wrapper.py b/boolformer/model/model_wrapper.py
index f2ed249..61db4a7 100755
--- a/boolformer/model/model_wrapper.py
+++ b/boolformer/model/model_wrapper.py
@@ -113,7 +113,7 @@ def fit(self,
                 #pred_tree.to_arbitrary_fan_in()
                 pred_tree.simplify()
                 pred_trees.append(pred_tree)
-                pred = pred_tree(inputs[0]).flatten()
+                pred = pred_tree(inputs[problem_idx]).flatten()
                 true = outputs[problem_idx]
                 error = 1.-sum(pred==true)/len(pred)
                 errors.append(error)