Skip to content

Commit

Permalink
fix: torch dependency start error (infiniflow#2777)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?

when use slim image, remove ```torch``` denpendency.

### Type of change

- [✓] Bug Fix (non-breaking change which fixes an issue)

---------

Co-authored-by: chongchuanbing <chongchuanbing@gmail.com>
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
  • Loading branch information
3 people authored Oct 10, 2024
1 parent 93bb243 commit be6f821
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions deepdoc/parser/pdf_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,12 @@ def __init__(self):

self.updown_cnt_mdl = xgb.Booster()
if not LIGHTEN:
import torch
if torch.cuda.is_available():
self.updown_cnt_mdl.set_param({"device": "cuda"})
try:
import torch
if torch.cuda.is_available():
self.updown_cnt_mdl.set_param({"device": "cuda"})
except Exception as e:
logging.error(str(e))
try:
model_dir = os.path.join(
get_project_base_directory(),
Expand Down

0 comments on commit be6f821

Please sign in to comment.