Bladeren bron

fix: add conditional import for torch and torch_npu in config_reader.py

myhloli 5 maanden geleden
bovenliggende
commit
bd5252d946
1 gewijzigde bestanden met toevoegingen van 7 en 4 verwijderingen
  1. 7 4
      mineru/utils/config_reader.py

+ 7 - 4
mineru/utils/config_reader.py

@@ -1,9 +1,15 @@
 # Copyright (c) Opendatalab. All rights reserved.
 import json
 import os
-
 from loguru import logger
 
+try:
+    import torch
+    import torch_npu
+except ImportError:
+    pass
+
+
 # 定义配置文件名常量
 CONFIG_FILE_NAME = os.getenv('MINERU_TOOLS_CONFIG_JSON', 'mineru.json')
 
@@ -71,15 +77,12 @@ def get_device():
     if device_mode is not None:
         return device_mode
     else:
-        import torch
-
         if torch.cuda.is_available():
             return "cuda"
         elif torch.backends.mps.is_available():
             return "mps"
         else:
             try:
-                import torch_npu
                 if torch_npu.npu.is_available():
                     return "npu"
             except Exception as e: