diff --git a/echopop/inversion/inversion_length_TS.py b/echopop/inversion/inversion_length_TS.py
index f45ab9dc..1f5b176c 100644
--- a/echopop/inversion/inversion_length_TS.py
+++ b/echopop/inversion/inversion_length_TS.py
@@ -315,7 +315,7 @@ def get_stratified_sigma_bs(
             df_length_counts["sigma_bs"] = 10.0 ** (df_length_counts["TS"] / 10.0)
 
             # Aggregate by stratum
-            self.sigma_bs_strata = df_length_counts.groupby(["stratum_ks"])[
+            self.sigma_bs_strata = df_length_counts.groupby(self.model_params["stratify_by"])[
                 ["length_count", "sigma_bs"]
             ].apply(lambda x: np.average(x.sigma_bs, weights=x.length_count))