diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py
index 164afd5..060f8e2 100644
--- a/evaluate_auc_v2.py
+++ b/evaluate_auc_v2.py
@@ -1138,6 +1138,7 @@ def evaluate_landmark_auc(
     disease_ids: Sequence[int],
     disease_chunk_size: int,
     score_mode: str,
+    dist_mode: str,
     horizons: np.ndarray,
     device: torch.device,
     model_target_mode: str,
@@ -1208,7 +1209,8 @@ def evaluate_landmark_auc(
                 death_token_ids=np.asarray(
                     landmark_dataset.death_token_ids, dtype=np.int64),
                 dist_mode=dist_mode,
-                model_death_idx=int(getattr(model, "death_idx", dataset.vocab_size - 1)),
+                model_death_idx=int(getattr(
+                    model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
             )
             nested = [_eval_token(t) for t in tqdm(
                 tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
@@ -1236,7 +1238,8 @@ def evaluate_landmark_auc(
                     np.asarray(landmark_dataset.death_token_ids,
                                dtype=np.int64),
                     dist_mode,
-                    int(getattr(model, "death_idx", dataset.vocab_size - 1)),
+                    int(getattr(
+                        model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
                 ),
             ) as ex:
                 nested = list(
@@ -1607,6 +1610,7 @@ def main() -> None:
         disease_ids=disease_ids,
         disease_chunk_size=disease_chunk_size,
         score_mode=score_mode,
+        dist_mode=dist_mode,
         horizons=horizons,
         device=device,
         model_target_mode=model_target_mode,