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

fix: Ocr AccleratorDevice

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 15:23:56 +00:00
parent 6bc1bd2ec4
commit 5c69081453
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docling/models/easyocr_model.py
View File

@ -47,8 +47,8 @@ class EasyOcrModel(BaseOcrModel):
# Enable easyocr GPU if running on CUDA, MPS # Enable easyocr GPU if running on CUDA, MPS
use_gpu = any( use_gpu = any(
filter( filter(
lambda x: str(x) in device, lambda x: str(x).lower() in device,
["cuda", "mps"], [AcceleratorDevice.CUDA.value, AcceleratorDevice.MPS.value],
) )
) )

2
docling/models/rapid_ocr_model.py
View File

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