From 4ff4688486fa21a2f6b2886f8786e46851100bf7 Mon Sep 17 00:00:00 2001 From: MyunghoonJin Date: Thu, 11 Apr 2024 17:15:10 +0900 Subject: [PATCH] fix cuda device error for tableformer/unitable --- src/openparse/tables/table_transformers/ml.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/openparse/tables/table_transformers/ml.py b/src/openparse/tables/table_transformers/ml.py index 51d5eaa..ad958af 100644 --- a/src/openparse/tables/table_transformers/ml.py +++ b/src/openparse/tables/table_transformers/ml.py @@ -186,7 +186,7 @@ def _table_outputs_to_objs( def find_table_bboxes( image: Image.Image, min_table_confidence: float ) -> List[_TableModelOutput]: - pixel_values = detection_transform(image).unsqueeze(0).to("cpu") + pixel_values = detection_transform(image).unsqueeze(0).to(device) with torch.no_grad(): outputs = detection_model(pixel_values) @@ -334,7 +334,7 @@ def get_table_content( len(structure_model.config.id2label): "no object", } - pixel_values_st = structure_transform(table_img).unsqueeze(0).to("cpu") + pixel_values_st = structure_transform(table_img).unsqueeze(0).to(device) with torch.no_grad(): outputs_st = structure_model(pixel_values_st)