diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py
index 57497ed..d1f33a7 100644
--- a/evaluate_event_free_survival.py
+++ b/evaluate_event_free_survival.py
@@ -291,18 +291,17 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
 def build_group_ablated_slice(
     batch: Dict[str, torch.Tensor],
     token_ids: Sequence[int],
-    row_start: int,
-    row_stop: int,
+    row_indices: torch.Tensor,
 ) -> Dict[str, torch.Tensor]:
     """Build one fixed-width ablated slice without rebuilding variable-length rows."""
     event_seq = batch["event_seq"]
 
     out: Dict[str, torch.Tensor] = {}
-    out["event_seq"] = event_seq[row_start:row_stop].clone()
-    out["time_seq"] = batch["time_seq"][row_start:row_stop]
-    out["readout_mask"] = batch["readout_mask"][row_start:row_stop].clone()
-    out["padding_mask"] = batch["padding_mask"][row_start:row_stop].bool().clone()
-    out["landmark_pos"] = batch["landmark_pos"][row_start:row_stop].clone()
+    out["event_seq"] = event_seq[row_indices].clone()
+    out["time_seq"] = batch["time_seq"][row_indices]
+    out["readout_mask"] = batch["readout_mask"][row_indices].clone()
+    out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
+    out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
 
     seq_len = int(event_seq.shape[1])
     positions = torch.arange(seq_len, device=event_seq.device)[None, :]
@@ -320,9 +319,9 @@ def build_group_ablated_slice(
     if not bool(has_valid.all().item()):
         empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
         out["event_seq"][empty_rows, 0] = CHECKUP_IDX
-        out["time_seq"][empty_rows, 0] = batch["t_query"][
-            row_start + empty_rows
-        ].to(dtype=out["time_seq"].dtype)
+        out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
+            dtype=out["time_seq"].dtype
+        )
         out["padding_mask"][empty_rows, 0] = True
         out["readout_mask"][empty_rows, 0] = True
         out["landmark_pos"][empty_rows] = 0
@@ -354,7 +353,7 @@ def build_group_ablated_slice(
         "row_idx",
     )
     for key in repeated_keys:
-        out[key] = batch[key][row_start:row_stop]
+        out[key] = batch[key][row_indices]
     return out
 
 
@@ -373,7 +372,6 @@ def iter_group_ablated_batches(
     max_batch_size: int,
 ):
     """Yield ablated chunks as soon as enough rows are available for a forward pass."""
-    batch_n = int(batch["event_seq"].shape[0])
     pending_batches: list[Dict[str, torch.Tensor]] = []
     pending_groups: list[str] = []
     pending_rows: list[int] = []
@@ -381,25 +379,28 @@ def iter_group_ablated_batches(
 
     for group in group_names:
         ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
-        if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
+        if ids.numel() == 0:
+            continue
+        active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten()
+        if active_rows.numel() == 0:
             continue
 
-        row_start = 0
-        while row_start < batch_n:
+        row_offset = 0
+        while row_offset < int(active_rows.numel()):
             capacity = int(max_batch_size) - pending_n
-            row_stop = min(batch_n, row_start + capacity)
+            row_stop = min(int(active_rows.numel()), row_offset + capacity)
+            row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
             chunk = build_group_ablated_slice(
                 batch=batch,
                 token_ids=organ_groups[group],
-                row_start=row_start,
-                row_stop=row_stop,
+                row_indices=row_indices,
             )
-            chunk_n = int(row_stop - row_start)
+            chunk_n = int(row_indices.numel())
             pending_batches.append(chunk)
             pending_groups.extend([group] * chunk_n)
-            pending_rows.extend(range(row_start, row_stop))
+            pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
             pending_n += chunk_n
-            row_start = row_stop
+            row_offset = row_stop
 
             if pending_n >= int(max_batch_size):
                 yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
@@ -774,6 +775,9 @@ def main() -> None:
             for group in group_names:
                 out[f"history_count__{group}"] = int(group_counts[group])
                 out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
+                if int(group_counts[group]) == 0:
+                    group_mortality_attr_prob[group][j] = 0.0
+                    group_mortality_attr_hazard[group][j] = 0.0
                 out[f"mortality_attribution_probability__{group}"] = float(
                     group_mortality_attr_prob[group][j]
                 )