diff --git a/miplearn/branching.py b/miplearn/branching.py
index 3163a80..7a3d0d0 100644
--- a/miplearn/branching.py
+++ b/miplearn/branching.py
@@ -5,61 +5,140 @@
 from . import Component
 from .transformers import PerVariableTransformer
 from abc import ABC, abstractmethod
+from sklearn.neighbors import KNeighborsRegressor
 import numpy as np
+from p_tqdm import p_map
 
 
+from tqdm.auto import tqdm
+from joblib import Parallel, delayed
+import multiprocessing
+
 class BranchPriorityComponent(Component):
     def __init__(self,
-                 initial_priority=None,
-                 collect_training_data=True):
-        self.priority = initial_priority
+                 node_limit=1_000,
+                ):
         self.transformer = PerVariableTransformer()
-        self.collect_training_data = collect_training_data
+        self.pending_instances = []
+        self.x_train = {}
+        self.y_train = {}
+        self.predictors = {}
+        self.node_limit = node_limit
     
     def before_solve(self, solver, instance, model):
         assert solver.is_persistent, "BranchPriorityComponent requires a persistent solver"
+        from gurobipy import GRB
         var_split = self.transformer.split_variables(instance, model)
         for category in var_split.keys():
+            if category not in self.predictors.keys():
+                continue
             var_index_pairs = var_split[category]
-            if self.priority is not None:
-                from gurobipy import GRB
-                for (i, (var, index)) in enumerate(var_index_pairs):
-                    gvar = solver.internal_solver._pyomo_var_to_solver_var_map[var[index]]
-                    gvar.setAttr(GRB.Attr.BranchPriority, int(self.priority[index]))        
+            for (i, (var, index)) in enumerate(var_index_pairs):
+                x = self._build_x(instance, var, index)
+                y = self.predictors[category].predict([x])[0][0]
+                gvar = solver.internal_solver._pyomo_var_to_solver_var_map[var[index]]
+                gvar.setAttr(GRB.Attr.BranchPriority, int(round(y)))
 
                     
     def after_solve(self, solver, instance, model):
-        if self.collect_training_data:
-            import subprocess, tempfile, os
-            src_dirname = os.path.dirname(os.path.realpath(__file__))
-            model_file = tempfile.NamedTemporaryFile(suffix=".lp")
+        self.pending_instances += [instance]
+    
+    def fit(self, solver, n_jobs=1):
+        def _process(instance):
+            # Create LP file
+            import subprocess, tempfile, os, sys
+            lp_file = tempfile.NamedTemporaryFile(suffix=".lp")
             priority_file = tempfile.NamedTemporaryFile()
-            solver.internal_solver.write(model_file.name)
+            model = instance.to_model()
+            model.write(lp_file.name)
+            
+            # Run Julia script
+            src_dirname = os.path.dirname(os.path.realpath(__file__))
+            priority_file = tempfile.NamedTemporaryFile(mode="r")
             subprocess.run(["julia",
                             "%s/scripts/branchpriority.jl" % src_dirname,
-                            model_file.name,
-                            priority_file.name],
+                            lp_file.name,
+                            priority_file.name,
+                            str(self.node_limit),
+                           ],
                            check=True,
-                           capture_output=True)
-            self._merge(np.genfromtxt(priority_file.name,
-                                      delimiter=',',
-                                      dtype=np.float64))
+                          )
+            
+            # Parse output
+            tokens = [line.strip().split(",") for line in priority_file.readlines()]
+            lp_varname_to_priority = {t[0]: int(t[1]) for t in tokens}
         
-    
-    def fit(self, solver):
-        pass
+            # Map priorities back to Pyomo variables
+            pyomo_var_to_priority = {}
+            from pyomo.core import Var
+            from pyomo.core.base.label import TextLabeler
+            labeler = TextLabeler()
+            symbol_map = list(model.solutions.symbol_map.values())[0]
+            
+            # Build x_train and y_train
+            comp = BranchPriorityComponent()
+            for var in model.component_objects(Var):
+                for index in var:
+                    category = instance.get_variable_category(var, index)
+                    if category is None:
+                        continue
+                    lp_varname = symbol_map.getSymbol(var[index], labeler)
+                    var_priority = lp_varname_to_priority[lp_varname]
+                    x = self._build_x(instance, var, index)
+                    y = np.array([var_priority])
+                    
+                    if category not in comp.x_train.keys():
+                        comp.x_train[category] = np.array([x])
+                        comp.y_train[category] = np.array([y])
+                    else:
+                        comp.x_train[category] = np.vstack([comp.x_train[category], x])
+                        comp.y_train[category] = np.vstack([comp.y_train[category], y])
+                
+            return comp
+        
+        
+        subcomponents = Parallel(n_jobs=n_jobs)(
+            delayed(_process)(instance)
+            for instance in tqdm(self.pending_instances, desc="Branch priority")
+        )
+        self.merge(subcomponents)
+        self.pending_instances.clear()
 
-    
+        # Retrain ML predictors
+        for category in self.x_train.keys():
+            x_train = self.x_train[category]
+            y_train = self.y_train[category]
+            self.predictors[category] = KNeighborsRegressor(n_neighbors=1)
+            self.predictors[category].fit(x_train, y_train)
+            
+        
+    def _build_x(self, instance, var, index):
+        instance_features = instance.get_instance_features()
+        var_features = instance.get_variable_features(var, index)
+        return np.hstack([instance_features, var_features])
+            
     def merge(self, other_components):
+        keys = set(self.x_train.keys())
         for comp in other_components:
-            if comp.priority is not None:
-                self._merge(comp.priority)
+            self.pending_instances += comp.pending_instances
+            keys = keys.union(set(comp.x_train.keys()))
             
+        # Merge x_train and y_train
+        for key in keys:
+            x_train_submatrices = [comp.x_train[key]
+                                   for comp in other_components
+                                   if key in comp.x_train.keys()]
+            y_train_submatrices = [comp.y_train[key]
+                                   for comp in other_components
+                                   if key in comp.y_train.keys()]
+            if key in self.x_train.keys():
+                x_train_submatrices += [self.x_train[key]]
+                y_train_submatrices += [self.y_train[key]]
+            self.x_train[key] = np.vstack(x_train_submatrices)
+            self.y_train[key] = np.vstack(y_train_submatrices)
             
-    def _merge(self, priority):
-        assert isinstance(priority, np.ndarray)
-        if self.priority is None:
-            self.priority = priority
-        else:
-            assert self.priority.shape == priority.shape
-            self.priority += priority
\ No newline at end of file
+        # Merge trained ML predictors
+        for comp in other_components:
+            for key in comp.predictors.keys():
+                if key not in self.predictors.keys():
+                    self.predictors[key] = comp.predictors[key]
\ No newline at end of file
diff --git a/miplearn/scripts/branchpriority.jl b/miplearn/scripts/branchpriority.jl
index 69aadc8..ff85ae2 100644
--- a/miplearn/scripts/branchpriority.jl
+++ b/miplearn/scripts/branchpriority.jl
@@ -3,6 +3,7 @@ using TinyBnB, CPLEXW, Printf
 
 instance_name = ARGS[1]
 output_filename = ARGS[2]
+node_limit = parse(Int, ARGS[3])
 
 mip = open_mip(instance_name)
 n_vars = CPXgetnumcols(mip.cplex_env[1], mip.cplex_lp[1])
@@ -34,12 +35,12 @@ function full_strong_branching_track(node::Node, progress::Progress)::TinyBnB.Va
     max_score, max_offset = findmax(scores)
     selected_var = node.fractional_variables[max_offset]
     
-    if rates_up[max_offset] < 1e6
+    if abs(rates_up[max_offset]) < 1e6
         pseudocost_count_up[selected_var.index] += 1
         pseudocost_sum_up[selected_var.index] += rates_up[max_offset]
     end
     
-    if rates_down[max_offset] < 1e6
+    if abs(rates_down[max_offset]) < 1e6
         pseudocost_count_down[selected_var.index] += 1
         pseudocost_sum_down[selected_var.index] += rates_down[max_offset]
     end
@@ -48,7 +49,7 @@ function full_strong_branching_track(node::Node, progress::Progress)::TinyBnB.Va
 end
 
 branch_and_bound(mip,
-                 node_limit = 1000,
+                 node_limit = node_limit,
                  branch_rule = full_strong_branching_track,
                  node_rule = best_bound,
                  print_interval = 100)
@@ -59,9 +60,7 @@ priority = [(pseudocost_count_up[v] == 0 || pseudocost_count_down[v] == 0) ? 0 :
              for v in 1:n_vars];
 
 open(output_filename, "w") do file
-    for v in 1:n_vars
-        v == 1 || write(file, ",")
-        write(file, @sprintf("%.0f", priority[v]))
+    for var in mip.binary_variables
+        write(file, @sprintf("%s,%.0f\n", name(mip, var), priority[var.index]))
     end
-    write(file, "\n")
-end
\ No newline at end of file
+end
diff --git a/miplearn/tests/test_branching.py b/miplearn/tests/test_branching.py
new file mode 100644
index 0000000..b39c710
--- /dev/null
+++ b/miplearn/tests/test_branching.py
@@ -0,0 +1,49 @@
+# MIPLearn, an extensible framework for Learning-Enhanced Mixed-Integer Optimization
+# Copyright (C) 2019-2020 Argonne National Laboratory. All rights reserved.
+# Written by Alinson S. Xavier <axavier@anl.gov>
+
+from miplearn import BranchPriorityComponent, LearningSolver
+from miplearn.problems.knapsack import MultiKnapsackInstance
+import numpy as np
+import tempfile
+
+def _get_instances():
+    return [
+        MultiKnapsackInstance(
+            weights=np.array([[23., 26., 20., 18.]]),
+            prices=np.array([505., 352., 458., 220.]),
+            capacities=np.array([67.])
+        ),
+    ] * 2
+
+
+def test_branching():
+    instances = _get_instances()
+    component = BranchPriorityComponent()
+    for instance in instances:
+        component.after_solve(None, instance, None)
+    component.fit(None)
+    for key in [0, 1, 2, 3]:
+        assert key in component.x_train.keys()
+        assert key in component.y_train.keys()
+        assert component.x_train[key].shape == (2,  9)
+        assert component.y_train[key].shape == (2,  1)
+        
+        
+def test_branch_priority_save_load():
+    state_file = tempfile.NamedTemporaryFile(mode="r")
+    solver = LearningSolver(components={"branch-priority": BranchPriorityComponent()})
+    solver.parallel_solve(_get_instances(), n_jobs=2)
+    solver.fit()
+    comp = solver.components["branch-priority"]
+    assert comp.x_train[0].shape == (2, 9)
+    assert comp.y_train[0].shape == (2, 1)
+    assert 0 in comp.predictors.keys()
+    solver.save_state(state_file.name)
+    
+    solver = LearningSolver(components={"branch-priority": BranchPriorityComponent()})
+    solver.load_state(state_file.name)
+    comp = solver.components["branch-priority"]
+    assert comp.x_train[0].shape == (2, 9)
+    assert comp.y_train[0].shape == (2, 1)
+    assert 0 in comp.predictors.keys()