Merge pull request #60 from ihmeuw-msca/bugfix/ensemble-infer-shape

Add `_infer_shape` function to the ensemble class

Author: zhengp0

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
         )