瀏覽代碼

fix: disable cuDNN for MACA device in common.py to improve compatibility

myhloli 1 周之前
父節點
當前提交
34c46cb83d
共有 1 個文件被更改,包括 5 次插入0 次删除
  1. 5 0
      mineru/cli/common.py

+ 5 - 0
mineru/cli/common.py

@@ -18,6 +18,11 @@ from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
 from mineru.backend.vlm.vlm_analyze import aio_doc_analyze as aio_vlm_doc_analyze
 from mineru.utils.pdf_page_id import get_end_page_id
 
+if os.getenv("MINERU_LMDEPLOY_DEVICE", "") == "maca":
+    import torch
+    torch.backends.cudnn.enabled = False
+
+
 pdf_suffixes = ["pdf"]
 image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]