diff --git a/evaluate_auc.py b/evaluate_auc.py
index 0c9265a..da867c1 100644
--- a/evaluate_auc.py
+++ b/evaluate_auc.py
@@ -632,15 +632,19 @@ def compute_logits_for_disease_chunk(
         device.type == "cuda" and use_amp) else torch.float32
     weight = model.risk_head.weight[disease_ids].detach().to(
         device=device, dtype=compute_dtype)
-    bias = model.risk_head.bias[disease_ids].detach().to(
-        device=device, dtype=compute_dtype)
+    bias = None
+    if model.risk_head.bias is not None:
+        bias = model.risk_head.bias[disease_ids].detach().to(
+            device=device, dtype=compute_dtype)
 
     parts: List[np.ndarray] = []
     for start in tqdm(range(0, n, logit_batch_size), desc="Risk-head projection", leave=False, dynamic_ncols=True):
         end = min(start + logit_batch_size, n)
         h = torch.from_numpy(hidden_all[start:end]).to(
             device=device, dtype=compute_dtype, non_blocking=True)
-        logits = torch.matmul(h, weight.t()) + bias
+        logits = torch.matmul(h, weight.t())
+        if bias is not None:
+            logits = logits + bias
         parts.append(logits.float().cpu().numpy().astype(
             np.float32, copy=False))
         del h, logits
diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py
index acdddec..8aef743 100644
--- a/evaluate_auc_v2.py
+++ b/evaluate_auc_v2.py
@@ -839,8 +839,10 @@ def project_distribution_chunk(
         device.type == "cuda" and use_amp) else torch.float32
     weight = model.risk_head.weight[disease_ids].detach().to(
         device=device, dtype=compute_dtype)
-    bias = model.risk_head.bias[disease_ids].detach().to(
-        device=device, dtype=compute_dtype)
+    bias = None
+    if model.risk_head.bias is not None:
+        bias = model.risk_head.bias[disease_ids].detach().to(
+            device=device, dtype=compute_dtype)
     rho_weight = None
     rho_bias = None
     death_rho_weight = None
@@ -868,7 +870,9 @@ def project_distribution_chunk(
         end = min(start + logit_batch_size, n)
         h = torch.from_numpy(hidden_all[start:end]).to(
             device=device, dtype=compute_dtype, non_blocking=True)
-        logits = torch.matmul(h, weight.t()) + bias
+        logits = torch.matmul(h, weight.t())
+        if bias is not None:
+            logits = logits + bias
         rho = None
         if dist_mode == "weibull":
             assert rho_weight is not None and rho_bias is not None