1

pip install onnxruntime-gpu==1.18.0

1
2
3
4
5
6
7
8
9

import onnxruntime
model_path = 'test.onnx'
CUDAExecutionProvider = {
                "cudnn_conv_algo_search": "DEFAULT",
                "cudnn_conv_use_max_workspace": '1',
                "device_id": '1' 
            }
providers=[("CUDAExecutionProvider", CUDAExecutionProvider)]
session = onnxruntime.InferenceSession(model_path, providers)

1

docker --gpus 'devices=1'