From 00de2ff4a94b2ec64053287707c7690fbf528775 Mon Sep 17 00:00:00 2001
From: Jiarui Li <jiarui_li@outlook.com>
Date: Sat, 27 Jun 2026 12:08:19 +0800
Subject: [PATCH] Fix event-free survival inference tensor use

---
 evaluate_event_free_survival.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py
index 292d09b..a03f055 100644
--- a/evaluate_event_free_survival.py
+++ b/evaluate_event_free_survival.py
@@ -284,7 +284,7 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
     }
 
 
-@torch.inference_mode()
+@torch.no_grad()
 def infer_landmark_hidden(
     *,
     model: DeepHealth,