diff --git a/inst/cohort-scratch.R b/inst/cohort-scratch.R
index 61d3b5b27..0131e8bf3 100644
--- a/inst/cohort-scratch.R
+++ b/inst/cohort-scratch.R
@@ -49,7 +49,7 @@ summary(fit_direct_weighted)
 
 lognormal <- brms::lognormal()
 
-primarycensored_lognormal_uniform_lpmf <- brms::custom_family(
+primarycensored_lognormal_uniform <- brms::custom_family(
   "primarycensored_lognormal_uniform",
   dpars = lognormal$dpar,
   links = c(lognormal$link, lognormal$link_sigma),
@@ -92,7 +92,7 @@ stanvars_all <- stanvars_functions + stanvars_data
 
 stancode <- brms::make_stancode(
   formula = d | weights(n) + vreal(q) ~ 1,
-  family = primarycensored_lognormal_uniform_lpmf,
+  family = primarycensored_lognormal_uniform,
   data = data,
   stanvars = stanvars_all,
 )
@@ -101,7 +101,7 @@ model <- rstan::stan_model(model_code = stancode)
 
 fit_pcd <- brms::brm(
   formula = d | weights(n) + vreal(q) ~ 1,
-  family = primarycensored_lognormal_uniform_lpmf,
+  family = primarycensored_lognormal_uniform,
   data = data,
   stanvars = stanvars_all,
   backend = "cmdstanr"
diff --git a/inst/stan/cohort_model/primarycensored-edit.stan b/inst/stan/cohort_model/primarycensored-edit.stan
index 0b2899b70..d0a77774e 100644
--- a/inst/stan/cohort_model/primarycensored-edit.stan
+++ b/inst/stan/cohort_model/primarycensored-edit.stan
@@ -2,10 +2,12 @@
 // Edited to work with temporary brms function
 
 real primarycensored_lognormal_uniform_lpmf(data int d, real mu, real sigma, real q, data real pwindow) {
+  int dist_id = 1; // lognormal
   array[2] real params = {mu, sigma};
-  array[0] real primary_params;
   int d_upper = d + 1;
-  return primarycensored_lpmf(d | 1, params, pwindow, d_upper, positive_infinity(), 1, primary_params);
+  int primary_id = 1; // Uniform
+  array[0] real primary_params;
+  return primarycensored_lpmf(d | dist_id, params, pwindow, d_upper, positive_infinity(), primary_id, primary_params);
 }
 
 /**