diff --git a/benchmark/benchmark.py b/benchmark/benchmark.py
index 8941705..e7f971a 100755
--- a/benchmark/benchmark.py
+++ b/benchmark/benchmark.py
@@ -64,6 +64,7 @@ def train():
         },
     )
     solver.parallel_solve(train_instances, n_jobs=10)
+    solver.fit(n_jobs=10)
     solver.save_state("%s/training_data.bin" % basepath)
     save(train_instances, "%s/train_instances.bin" % basepath)
     save(test_instances, "%s/test_instances.bin" % basepath)
@@ -103,7 +104,6 @@ def test_ml():
     test_instances = load("%s/test_instances.bin" % basepath)
     benchmark = BenchmarkRunner(solvers)
     benchmark.load_state("%s/training_data.bin" % basepath)
-    benchmark.fit()
     benchmark.load_results("%s/benchmark_baseline.csv" % basepath)
     benchmark.parallel_solve(test_instances, n_jobs=10)
     benchmark.save_results("%s/benchmark_ml.csv" % basepath)    
diff --git a/miplearn/solvers.py b/miplearn/solvers.py
index 47144c8..d891e18 100644
--- a/miplearn/solvers.py
+++ b/miplearn/solvers.py
@@ -9,11 +9,8 @@ import pyomo.environ as pe
 import numpy as np
 from copy import deepcopy
 import pickle
-from tqdm import tqdm
-from joblib import Parallel, delayed
 from scipy.stats import randint
-import multiprocessing
-
+from p_tqdm import p_map
 
 def _gurobi_factory():
     solver = pe.SolverFactory('gurobi_persistent')
@@ -29,7 +26,9 @@ class LearningSolver:
     """
 
     def __init__(self,
-                 threads=4,
+                 threads=None,
+                 time_limit=None,
+                 gap_limit=None,
                  internal_solver_factory=_gurobi_factory,
                  components=None,
                  mode=None):
@@ -37,6 +36,9 @@ class LearningSolver:
         self.internal_solver = None
         self.components = components
         self.internal_solver_factory = internal_solver_factory
+        self.threads = threads
+        self.time_limit = time_limit
+        self.gap_limit = gap_limit
         
         if self.components is not None:
             assert isinstance(self.components, dict)
@@ -54,6 +56,12 @@ class LearningSolver:
     def _create_solver(self):
         self.internal_solver = self.internal_solver_factory()
         self.is_persistent = hasattr(self.internal_solver, "set_instance")
+        if self.threads is not None:
+            self.internal_solver.options["Threads"] = self.threads
+        if self.time_limit is not None:
+            self.internal_solver.options["TimeLimit"] = self.time_limit
+        if self.gap_limit is not None:
+            self.internal_solver.options["MIPGap"] = self.gap_limit
         
     def solve(self, instance, tee=False):
         model = instance.to_model()
@@ -93,11 +101,7 @@ class LearningSolver:
                 solver.components = {}
             return solver, results
 
-        solver_result_pairs = Parallel(n_jobs=n_jobs)(
-            delayed(_process)(instance)
-            for instance in tqdm(instances, desc=label, ncols=80)
-        )
-        
+        solver_result_pairs = p_map(_process, instances, num_cpus=n_jobs, desc=label)
         subsolvers = [p[0] for p in solver_result_pairs]
         results = [p[1] for p in solver_result_pairs]
         
@@ -109,9 +113,9 @@ class LearningSolver:
         
         return results
 
-    def fit(self):
+    def fit(self, n_jobs=1):
         for component in self.components.values():
-            component.fit(self)
+            component.fit(self, n_jobs=n_jobs)
             
     def save_state(self, filename):
         with open(filename, "wb") as file:
diff --git a/miplearn/tests/test_solver.py b/miplearn/tests/test_solver.py
index fc7c532..e5c5ebe 100644
--- a/miplearn/tests/test_solver.py
+++ b/miplearn/tests/test_solver.py
@@ -56,7 +56,7 @@ def test_parallel_solve():
 def test_solver_random_branch_priority():
     instance = _get_instance()
     components = {
-        "warm-start": BranchPriorityComponent(initial_priority=np.array([1., 2., 3., 4.])),
+        "warm-start": BranchPriorityComponent(),
     }
     solver = LearningSolver(components=components)
     solver.solve(instance)
diff --git a/requirements.txt b/requirements.txt
index f3c7116..a7e0f14 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,15 @@
 docopt
+matplotlib
 mkdocs
 mkdocs-cinder
 networkx
 numpy
+p_tqdm
 pandas
 pyomo
 pytest
 pytest-watch
 python-markdown-math
+seaborn
 sklearn
 tqdm