Refactor evaluation device handling in AUC scripts for improved clarity and flexibility

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