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,