diff --git a/compiler/compile-config/compile-config-2.02-learning_inference b/compiler/compile-config/compile-config-2.02-learning_inference
index ca1152037..74fdfda3e 100755
--- a/compiler/compile-config/compile-config-2.02-learning_inference
+++ b/compiler/compile-config/compile-config-2.02-learning_inference
@@ -196,12 +196,7 @@ if .deepdive_.execution.processes | has("process/grounding/combine_factorgraph")
             for pid in $DEEPDIVE_FACTORGRAPH_SHARDS; do
                 cd \"$DEEPDIVE_RESULTS_DIR\"/$pid
                 cat inference_result.out.text |
-                # restoring shard ID to the vids
-                python3 -c "
-import sys
-for line in sys.stdin:
-  vid, cid, prob = line.split(); print('\t'.join([str(int(vid) + ($SHARD_BASE << 48)), cid, prob]))
-" |
+                pid=$pid deepdive env restore_partitioned_vids |
                 DEEPDIVE_LOAD_FORMAT=tsv \\
                 deepdive load \(deepdiveInferenceResultVariablesTable | @sh) /dev/stdin
             done
diff --git a/inference/restore_partitioned_vids b/inference/restore_partitioned_vids
new file mode 100755
index 000000000..519a7f63f
--- /dev/null
+++ b/inference/restore_partitioned_vids
@@ -0,0 +1,7 @@
+#!/usr/bin/perl -w
+use strict;
+my $SHARD_BASE = $ENV{pid} << 48;
+while (<>) {
+    my ($vid, @rest) = split(" ");
+    print join("\t", ($SHARD_BASE | $vid), @rest), "\n";
+}
diff --git a/stage.sh b/stage.sh
index 137dd16a4..be3accf22 100755
--- a/stage.sh
+++ b/stage.sh
@@ -142,6 +142,7 @@ generate-wrapper-for-libdirs          "$STAGE_DIR"/util/sampler-$cmd \
 done
 stage inference/deepdive-model                                    util/
 stage inference/run-sampler                                       util/
+stage inference/restore_partitioned_vids                          util/
 
 # Stanford CoreNLP utilities
 stage util/nlp/deepdive-corenlp                                   util/