Browse Source

fix: enhance CUDA and NPU availability checks in utils.py

myhloli 1 week ago
parent
commit
1c530f64f5
1 changed files with 7 additions and 5 deletions
  1. 7 5
      mineru/backend/vlm/utils.py

+ 7 - 5
mineru/backend/vlm/utils.py

@@ -12,14 +12,16 @@ def enable_custom_logits_processors() -> bool:
     import torch
     from vllm import __version__ as vllm_version
 
-    if not torch.cuda.is_available():
+    if torch.cuda.is_available():
+        major, minor = torch.cuda.get_device_capability()
+        # 正确计算Compute Capability
+        compute_capability = f"{major}.{minor}"
+    elif torch.npu.is_available():
+        compute_capability = "8.0"
+    else:
         logger.info("CUDA not available, disabling custom_logits_processors")
         return False
 
-    major, minor = torch.cuda.get_device_capability()
-    # 正确计算Compute Capability
-    compute_capability = f"{major}.{minor}"
-
     # 安全地处理环境变量
     vllm_use_v1_str = os.getenv('VLLM_USE_V1', "1")
     if vllm_use_v1_str.isdigit():