PDF_Translate / app.py
qianluxi's picture
Update app.py
24dd54c verified
Raw
History Blame Contribute Delete
520 Bytes
import os
import onnxruntime as ort
# 强制 ONNX 使用 GPU(如果可用)
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
print("Available providers:", ort.get_available_providers())
# 创建推理会话时指定 providers
# model_path = "your_model.onnx"
# session = ort.InferenceSession(model_path, providers=providers)
# 打印当前设备
if 'CUDAExecutionProvider' in ort.get_available_providers():
print("✅ Running on GPU")
else:
print("⚠️ Running on CPU (no GPU allocated)")