diff --git a/src/python/miplearn/benchmark.py b/src/python/miplearn/benchmark.py
index b71c808..aae748c 100644
--- a/src/python/miplearn/benchmark.py
+++ b/src/python/miplearn/benchmark.py
@@ -28,8 +28,7 @@ class BenchmarkRunner:
         for (name, solver) in self.solvers.items():
             results = solver.parallel_solve(instances,
                                             n_jobs=n_jobs,
-                                            label="Solve (%s)" % name,
-                                            collect_training_data=False)
+                                            label="Solve (%s)" % name)
             for i in range(len(instances)):
                 self._push_result(results[i],
                                   solver=solver,
diff --git a/src/python/miplearn/tests/test_solver.py b/src/python/miplearn/tests/test_solver.py
index 5f9846a..8b92a1d 100644
--- a/src/python/miplearn/tests/test_solver.py
+++ b/src/python/miplearn/tests/test_solver.py
@@ -2,6 +2,9 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import pickle
+import tempfile
+
 from miplearn import LearningSolver, BranchPriorityComponent
 from miplearn.problems.knapsack import KnapsackInstance
 from miplearn.solvers import GurobiSolver