Commit a5b35315 authored by fengyf1's avatar fengyf1
Browse files

Initial commit

parent 9292fe06
Pipeline #3372 canceled with stages
...@@ -42,12 +42,12 @@ def load_model(model_name): ...@@ -42,12 +42,12 @@ def load_model(model_name):
options.enable_mem_reuse = False options.enable_mem_reuse = False
# Increase the number for faster inference and more memory consumption # Increase the number for faster inference and more memory consumption
options.intra_op_num_threads = 1 options.intra_op_num_threads = 1
cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',} rocm_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
session = ort.InferenceSession( session = ort.InferenceSession(
model_name, model_name,
sess_options=options, sess_options=options,
providers=[('CUDAExecutionProvider', cuda_provider_options)] providers=[('ROCMExecutionProvider', rocm_provider_options)]
) )
return session return session
......
onnx==1.14.0 onnx==1.14.0
onnxruntime-gpu==1.14.0 #onnxruntime-gpu==1.14.0
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment