Added cuda_is_available (#4725)

### What problem does this PR solve?

Added cuda_is_available

### Type of change

- [x] Refactoring
This commit is contained in:
Zhichang Yu
2025-02-05 18:01:23 +08:00
committed by GitHub
parent 283d036cba
commit 3411d0a2ce
4 changed files with 24 additions and 7 deletions

View File

@ -21,7 +21,6 @@ import numpy as np
import cv2
from copy import deepcopy
import torch
import onnxruntime as ort
from huggingface_hub import snapshot_download
@ -60,11 +59,21 @@ class Recognizer(object):
if not os.path.exists(model_file_path):
raise ValueError("not find model file path {}".format(
model_file_path))
def cuda_is_available():
try:
import torch
if torch.cuda.is_available():
return True
except Exception:
return False
return False
# https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
# Shrink GPU memory after execution
self.run_options = ort.RunOptions()
if torch.cuda.is_available():
if cuda_is_available():
options = ort.SessionOptions()
options.enable_cpu_mem_arena = False
cuda_provider_options = {