modified: estimators.py

autonlab · Jun 28, 2022 · 9e95ad1 · 9e95ad1
1 parent cc0f00e
commit 9e95ad1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/auton_survival/estimators.py b/auton_survival/estimators.py
@@ -228,6 +228,9 @@ def _predict_dcph(model, features, times):
   if isinstance(times, float) or isinstance(times, int):
     times = [float(times)]
 
+  if isinstance(times, np.ndarray):
+    times = times.ravel().tolist()
+
   return model.predict_survival(x=features.values, t=times)
 
 def _fit_cph(features, outcomes, val_data, random_seed, **hyperparams):
@@ -378,7 +381,7 @@ def _fit_dsm(features, outcomes, val_data, random_seed, **hyperparams):
 
   k = hyperparams.get("k", 3)
   layers = hyperparams.get("layers", [100])
-  epochs = hyperparams.get("iters", 10)
+  epochs = hyperparams.get("iters", 50)
   distribution = hyperparams.get("distribution", "Weibull")
   temperature = hyperparams.get("temperature", 1.0)
   lr = hyperparams.get("learning_rate", 1e-3)