(fix): Ensure survival tensors are moved to CPU before converting to …

…NumPy arrays.
arnor-sigurdsson · Jan 23, 2025 · 655ba6f · 655ba6f
1 parent 435e55b
commit 655ba6f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/eir/predict_modules/predict_output_modules/predict_survival_output.py b/eir/predict_modules/predict_output_modules/predict_survival_output.py
@@ -122,7 +122,7 @@ def predict_survival_wrapper_with_labels(
             )
 
         else:
-            risk_scores = model_outputs.numpy()
+            risk_scores = model_outputs.cpu().numpy()
 
             baseline_hazard = output_object.baseline_hazard
             unique_times = output_object.baseline_unique_times

diff --git a/eir/train_utils/evaluation_modules/evaluation_output_survival.py b/eir/train_utils/evaluation_modules/evaluation_output_survival.py
@@ -109,7 +109,7 @@ def save_survival_evaluation_results_wrapper(
             )
 
         else:
-            risk_scores = model_outputs.numpy()
+            risk_scores = model_outputs.cpu().numpy()
 
             unique_times, baseline_hazard = estimate_baseline_hazard(
                 times=times,