diff --git a/src/mrtool/core/model.py b/src/mrtool/core/model.py
index ca94e4c..df8bed3 100644
--- a/src/mrtool/core/model.py
+++ b/src/mrtool/core/model.py
@@ -523,6 +523,7 @@ def __init__(
 
         self.weights = np.ones(self.num_sub_models) / self.num_sub_models
 
+    def _infer_shape(self) -> None:
         # inherent the dimension variable
         self.num_x_vars = self.sub_models[0].num_x_vars
         self.num_z_vars = self.sub_models[0].num_z_vars
@@ -542,6 +543,8 @@ def fit_model(
         for sub_model in self.sub_models:
             sub_model.fit_model(**fit_options)
 
+        self._infer_shape()
+
         self.score_model(
             scores_weights=scores_weights, slopes=slopes, quantiles=quantiles
         )