From fb6cb8b04841a9ff5fd09896e6e122f8d1555ea3 Mon Sep 17 00:00:00 2001
From: ciaran <87439103+CiaranYoung@users.noreply.github.com>
Date: Mon, 4 Nov 2024 15:17:27 +0800
Subject: [PATCH] Update pdf_extract_kit.py

Modify line 397 to ensure compatibility with CPU execution, addressing the issue where specifying 'cpu' in config.json still results in a ValueError for expecting a cuda device but getting 'cpu' during demo execution.
---
 magic_pdf/model/pdf_extract_kit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py
index ec67e3ed..d6f07550 100644
--- a/magic_pdf/model/pdf_extract_kit.py
+++ b/magic_pdf/model/pdf_extract_kit.py
@@ -393,7 +393,7 @@ def __call__(self, image):
             elif int(res['category_id']) in [5]:
                 table_res_list.append(res)
 
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() and self.device != 'cpu':
             properties = torch.cuda.get_device_properties(self.device)
             total_memory = properties.total_memory / (1024 ** 3)  # 将字节转换为 GB
             if total_memory <= 10: