mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
Fixed GPU detection on CPU only environment (#4711)
### What problem does this PR solve? Fixed GPU detection on CPU only environment. Close #4692 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
@ -21,7 +21,7 @@ import numpy as np
|
||||
import cv2
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
import torch
|
||||
import onnxruntime as ort
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
@ -64,7 +64,7 @@ class Recognizer(object):
|
||||
# Shrink GPU memory after execution
|
||||
self.run_options = ort.RunOptions()
|
||||
|
||||
if ort.get_device() == "GPU":
|
||||
if torch.cuda.is_available():
|
||||
options = ort.SessionOptions()
|
||||
options.enable_cpu_mem_arena = False
|
||||
cuda_provider_options = {
|
||||
|
||||
Reference in New Issue
Block a user