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:
Zhichang Yu
2025-02-05 12:02:43 +08:00
committed by GitHub
parent 7a7f98b1a9
commit e1526846da
4 changed files with 7 additions and 6 deletions

View File

@ -27,6 +27,7 @@ from . import operators
import math
import numpy as np
import cv2
import torch
import onnxruntime as ort
from .postprocess import build_post_process
@ -80,7 +81,7 @@ def load_model(model_dir, nm):
# https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
# Shrink GPU memory after execution
run_options = ort.RunOptions()
if ort.get_device() == "GPU":
if torch.cuda.is_available():
cuda_provider_options = {
"device_id": 0, # Use specific GPU
"gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory

View File

@ -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 = {