Skip to content

GitLab

Unverified Commit af3ec555 authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #1567 from myhloli/dev 

refactor(table): add device configuration for Unitable model
parents fd5427a6 e64d4fed
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py
View file @ af3ec555
...@@ -5,6 +5,8 @@ from loguru import logger ...@@ -5,6 +5,8 @@ from loguru import logger
from rapid_table import RapidTable, RapidTableInput from rapid_table import RapidTable, RapidTableInput
from rapid_table.main import ModelType from rapid_table.main import ModelType
from magic_pdf.libs.config_reader import get_device
class RapidTableModel(object): class RapidTableModel(object):
def __init__(self, ocr_engine, table_sub_model_name): def __init__(self, ocr_engine, table_sub_model_name):
...@@ -13,7 +15,7 @@ class RapidTableModel(object): ...@@ -13,7 +15,7 @@ class RapidTableModel(object):
input_args = RapidTableInput() input_args = RapidTableInput()
elif table_sub_model_name in sub_model_list: elif table_sub_model_name in sub_model_list:
if torch.cuda.is_available() and table_sub_model_name == "unitable": if torch.cuda.is_available() and table_sub_model_name == "unitable":
input_args = RapidTableInput(model_type=table_sub_model_name, use_cuda=True) input_args = RapidTableInput(model_type=table_sub_model_name, use_cuda=True, device=get_device())
else: else:
input_args = RapidTableInput(model_type=table_sub_model_name) input_args = RapidTableInput(model_type=table_sub_model_name)
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment