fix broadcasting

sktime · Apr 17, 2024 · c7142c6 · c7142c6
1 parent f169b5f
commit c7142c6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/skpro/survival/aft/_aft_lifelines_weibull.py b/skpro/survival/aft/_aft_lifelines_weibull.py
@@ -3,6 +3,8 @@
 
 __author__ = ["fkiraly"]
 
+import numpy as np
+
 from skpro.distributions.weibull import Weibull
 from skpro.survival.adapters.lifelines import _LifelinesAdapter
 from skpro.survival.base import BaseSurvReg
@@ -181,6 +183,8 @@ def _predict_proba(self, X):
         ll_pred_proba = lifelines_est._prep_inputs_for_prediction_and_return_scores
 
         scale, k = ll_pred_proba(df, ancillary)
+        scale = np.expand_dims(scale, axis=1)
+        k = np.expand_dims(k, axis=1)
 
         dist = Weibull(scale=scale, k=k, index=X.index, columns=self._y_cols)
         return dist