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)")