Update dcph_utilities.py

autonlab · Jul 10, 2022 · 0b26fb6 · 0b26fb6
1 parent f86d654
commit 0b26fb6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/auton_survival/models/cph/dcph_utilities.py b/auton_survival/models/cph/dcph_utilities.py
@@ -143,7 +143,7 @@ def predict_survival(model, x, t=None):
   if isinstance(t, (int, float)): t = [t]
 
   model, breslow_spline = model
-  lrisks = model(x).detach().numpy()
+  lrisks = model(x).detach().cpu().numpy()
 
   unique_times = breslow_spline.baseline_survival_.x