saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-22 22:01:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixing optimizer failure due to missing provider list (#10497)

Browse source 
Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>
This commit is contained in:
Boris Fomitchev 2022-03-31 11:05:49 -07:00 committed by GitHub
parent bfcd5bd4a2
commit eab7c0d5bf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/python/tools/transformers/optimizer.py
View file

@ -93,7 +93,9 @@ def optimize_by_onnxruntime(onnx_model_path: str,
providers=['CPUExecutionProvider'],
**kwargs)
else:
session = onnxruntime.InferenceSession(onnx_model_path, sess_options, **kwargs)
session = onnxruntime.InferenceSession(onnx_model_path, sess_options,
providers=['CUDAExecutionProvider'],
**kwargs)
assert 'CUDAExecutionProvider' in session.get_providers() # Make sure there is GPU
assert os.path.exists(optimized_model_path) and os.path.isfile(optimized_model_path)