Update ocr_server.py

magic_pdf/dict2md/ocr_server.py

@@ -57,7 +57,7 @@ def _load_model_processor(args):
     else:
         if args.dcu_id is not None:
             device_map = {'': f'cuda:{args.dcu_id}'}
-            os.environ['CUDA_VISIBLE_DEVICES'] = args.dcu_id
+            
             print('使用DCU推理:', f'cuda:{args.dcu_id}')
         else:
             device_map = 'auto'
@@ -261,6 +261,7 @@ if __name__ == "__main__":
     args = _get_args()
     config = configparser.ConfigParser()
     config.read(args.config_path)
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.dcu_id
     # host = config.get('server', 'ocr_host')
     host, port = config.get('server', 'ocr_server').split('://')[1].split(':')[0], int(
         config.get('server', 'ocr_server').split('://')[1].split(':')[1])