mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-27 04:24:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Do proper check to set the device in EasyOCR, RapidOCR.

Browse Source 
Signed-off-by: Nikos Livathinos <nli@zurich.ibm.com>

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Nikos Livathinos 2024-12-10 14:46:21 +00:00
parent ce82e23b66
commit 99ccb69a47
2 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
docling/models/easyocr_model.py
View File

@ -45,10 +45,13 @@ class EasyOcrModel(BaseOcrModel):
if self.options.use_gpu:
device = decide_device(accelerator_options.device)
# Enable easyocr GPU if running on CUDA, MPS
use_gpu = device in [
AcceleratorDevice.CUDA,
AcceleratorDevice.MPS,
]
use_gpu = any(
filter(
lambda x: str(x) in device,
[AcceleratorDevice.CUDA, AcceleratorDevice.MPS],
)
)
self.reader = easyocr.Reader(
lang_list=self.options.lang,
gpu=use_gpu,

2
docling/models/rapid_ocr_model.py
View File

@ -42,7 +42,7 @@ class RapidOcrModel(BaseOcrModel):
# Decide the accelerator devices
device = decide_device(accelerator_options.device)
use_cuda = device == AcceleratorDevice.CUDA
use_cuda = AcceleratorDevice.CUDA in device
use_dml = accelerator_options.device == AcceleratorDevice.AUTO
intra_op_num_threads = accelerator_options.num_threads