Update app.py

app.py

@@ -281,6 +281,11 @@ def train_survival_bundle(
         vars_sorted = var1[X_cols].sort_values(ascending=False)
         X_cols = list(vars_sorted.head(k).index)
 
+    # ---- Phase 2 (CRITICAL): refit FINAL imputer on FINAL X_cols ----
+    imp = SimpleImputer(strategy="median", keep_empty_features=True)
+    X_final = imp.fit_transform(df_surv_oh[X_cols])
+    df_surv_oh.loc[:, X_cols] = pd.DataFrame(X_final, columns=X_cols, index=df_surv_oh.index)
+
     # ---- Fit Cox (use stronger ridge) ----
     # For small n / many predictors, penalizer 0.1 is often too weak.
     # Typical stable values: 1.0 to 10.0