From 615f39505887f3b08491c4adc45735bfed03e929 Mon Sep 17 00:00:00 2001
From: Jiarui Li <jiarui_li@outlook.com>
Date: Sat, 20 Jun 2026 12:13:58 +0800
Subject: [PATCH] Refactor evaluation device handling in AUC scripts for
 improved clarity and flexibility

---
 evaluate_auc.py    | 16 ++++++++++++++--
 evaluate_auc_v2.py | 16 ++++++++++++++--
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/evaluate_auc.py b/evaluate_auc.py
index 6a3fbe1..c22250b 100644
--- a/evaluate_auc.py
+++ b/evaluate_auc.py
@@ -298,6 +298,17 @@ def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any
     return cfg.get(name, default)
 
 
+def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
+    """Resolve evaluation device without inheriting train_config.json device."""
+    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device_name)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            f"Requested device {device_name!r}, but CUDA is not available."
+        )
+    return device
+
+
 def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     total = train_ratio + val_ratio + test_ratio
     if not np.isclose(total, 1.0, atol=1e-6):
@@ -1266,6 +1277,8 @@ def main() -> None:
                         help="Inference batch size; overrides train_config.json.")
     parser.add_argument("--num_workers", type=int, default=None,
                         help="DataLoader workers; overrides train_config.json.")
+    parser.add_argument("--device", type=str, default=None,
+                        help="Evaluation device, e.g. cpu, cuda, cuda:1. Defaults to cuda if available, else cpu.")
     parser.add_argument("--num_workers_auc", type=int, default=None,
                         help="CPU processes for AUC computation.")
     parser.add_argument("--auc_task_chunk_size", type=int, default=None,
@@ -1323,8 +1336,7 @@ def main() -> None:
         "readout_name", "same_time_group_end" if target_mode == "uts" else "token")
     readout_reduce = cfg.get("readout_reduce", "mean")
 
-    device = torch.device(cfg.get("device", "cuda")
-                          if torch.cuda.is_available() else "cpu")
+    device = resolve_eval_device(args.device)
     if device.type == "cuda":
         torch.backends.cudnn.benchmark = True
 
diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py
index 507b98f..164afd5 100644
--- a/evaluate_auc_v2.py
+++ b/evaluate_auc_v2.py
@@ -51,6 +51,17 @@ def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: A
     return cfg.get(name, default)
 
 
+def resolve_eval_device(device_arg: Optional[str]) -> torch.device:
+    """Resolve evaluation device without inheriting train_config.json device."""
+    device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device_name)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            f"Requested device {device_name!r}, but CUDA is not available."
+        )
+    return device
+
+
 def parse_int_list(value: Any) -> Optional[List[int]]:
     if value is None:
         return None
@@ -1316,6 +1327,8 @@ def main() -> None:
 
     parser.add_argument("--batch_size", type=int, default=None)
     parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument("--device", type=str, default=None,
+                        help="Evaluation device, e.g. cpu, cuda, cuda:1. Defaults to cuda if available, else cpu.")
     parser.add_argument("--num_workers_auc", type=int, default=None)
     parser.add_argument("--auc_task_chunk_size", type=int, default=None)
     parser.add_argument("--disease_chunk_size", type=int, default=None)
@@ -1470,8 +1483,7 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
 
-    device = torch.device(cfg.get("device", "cuda")
-                          if torch.cuda.is_available() else "cpu")
+    device = resolve_eval_device(args.device)
     if device.type == "cuda":
         torch.backends.cudnn.benchmark = True