diff --git a/src/HHbbVV/postprocessing/TrainBDT.py b/src/HHbbVV/postprocessing/TrainBDT.py
index aac46672..f954e9a3 100644
--- a/src/HHbbVV/postprocessing/TrainBDT.py
+++ b/src/HHbbVV/postprocessing/TrainBDT.py
@@ -807,7 +807,7 @@ def do_inference(
                 # have to change model's feature names since we're passing in a dataframe
                 model.get_booster().feature_names = mcvars
                 preds = model.predict_proba(X)
-                preds = preds[:, :-1] if multiclass else preds[:, 1]
+                preds = preds if multiclass else preds[:, 1]
                 np.save(f"{model_dir}/inferences/{year}/preds_{jshift}.npy", preds)
 
             for jshift in jmsr_shifts:
@@ -816,7 +816,7 @@ def do_inference(
                 # have to change model's feature names since we're passing in a dataframe
                 model.get_booster().feature_names = mcvars
                 preds = model.predict_proba(X)
-                preds = preds[:, :-1] if multiclass else preds[:, 1]
+                preds = preds if multiclass else preds[:, 1]
                 np.save(f"{model_dir}/inferences/{year}/preds_{jshift}.npy", preds)
 
 
diff --git a/src/HHbbVV/postprocessing/bash_scripts/NonresTemplates.sh b/src/HHbbVV/postprocessing/bash_scripts/NonresTemplates.sh
index 0da9a979..d51b3875 100755
--- a/src/HHbbVV/postprocessing/bash_scripts/NonresTemplates.sh
+++ b/src/HHbbVV/postprocessing/bash_scripts/NonresTemplates.sh
@@ -9,14 +9,19 @@
 MAIN_DIR="../../.."
 data_dir="$MAIN_DIR/../data/skimmer/24Mar14UpdateData"
 bdt_preds_dir="$data_dir/24_04_05_k2v0_training_eqsig_vbf_vars_rm_deta/inferences"
+sig_samples=""
 TAG=""
 
 
-options=$(getopt -o "" --long "tag:" -- "$@")
+options=$(getopt -o "" --long "sample:,tag:" -- "$@")
 eval set -- "$options"
 
 while true; do
     case "$1" in
+        --sample)
+            shift
+            sig_samples="--sig-samples $1"
+            ;;
         --tag)
             shift
             TAG=$1
@@ -43,7 +48,7 @@ fi
 
 for year in 2016APV 2016 2017 2018
 do
-    python -u postprocessing.py --year $year --data-dir "$data_dir" --templates \
+    python -u postprocessing.py --year $year --data-dir "$data_dir" --templates $sig_samples \
     --bdt-preds-dir $bdt_preds_dir \
     --plot-dir "${MAIN_DIR}/plots/PostProcessing/$TAG" \
     --template-dir "templates/$TAG" --plot-shifts --nonres-regions "vbf"