Skip to content

GitLab

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

Merge pull request #853 from CiaranYoung/master 

Update pdf_extract_kit.py
parents db01bcdc fb6cb8b0
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
magic_pdf/model/pdf_extract_kit.py
View file @ 8ee1da82
...@@ -393,7 +393,7 @@ class CustomPEKModel: ...@@ -393,7 +393,7 @@ class CustomPEKModel:
elif int(res['category_id']) in [5]: elif int(res['category_id']) in [5]:
table_res_list.append(res) 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) properties = torch.cuda.get_device_properties(self.device)
total_memory = properties.total_memory / (1024 ** 3) # 将字节转换为 GB total_memory = properties.total_memory / (1024 ** 3) # 将字节转换为 GB
if total_memory <= 10: if total_memory <= 10:
......
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