Merge pull request #1567 from myhloli/dev

refactor(table): add device configuration for Unitable model
opendatalab · Jan 17, 2025 · af3ec55 · af3ec55
2 parents fd5427a + e64d4fe
commit af3ec55
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py b/magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py
@@ -5,6 +5,8 @@
 from rapid_table import RapidTable, RapidTableInput
 from rapid_table.main import ModelType
 
+from magic_pdf.libs.config_reader import get_device
+
 
 class RapidTableModel(object):
     def __init__(self, ocr_engine, table_sub_model_name):
@@ -13,7 +15,7 @@ def __init__(self, ocr_engine, table_sub_model_name):
             input_args = RapidTableInput()
         elif table_sub_model_name in  sub_model_list:
             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:
                 input_args = RapidTableInput(model_type=table_sub_model_name)
         else: