Update ocr_server.py

magic_pdf/dict2md/ocr_server.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree
 import configparser
 import copy
+import os
 import re
 import gc
 import time
@@ -44,7 +45,7 @@ def _get_args():
     parser.add_argument('--dcu_id', type=str, default='0', help='Specify the GPU ID to load the model onto.')
     parser.add_argument(
         '--config_path',
-        default='/home/practice/magic_pdf-main/magic_pdf/config.ini',
+        default='./magic_pdf/config.ini',
         )
     args = parser.parse_args()
     return args
@@ -56,6 +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'