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

Fix cuda provider crash on it (#16056)

Browse source
This commit is contained in:
Zhang Lei 2023-05-24 16:13:11 -07:00 committed by GitHub
parent 105f5f0f20
commit 63c9973b7a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/contrib_ops/cuda/transformers/beam_search.cc
View file

@ -27,6 +27,7 @@ ONNX_OPERATOR_KERNEL_EX(
.InputMemoryType(OrtMemTypeCPUInput, 9) // 'attention_mask' needs to be on CPU
.InputMemoryType(OrtMemTypeCPUInput, 10) // 'decoder_input_ids' needs to be on CPU
.OutputMemoryType(OrtMemTypeCPUOutput, 0) // 'sequences' output on CPU
.InputMemoryType(OrtMemTypeCPUInput, 11) // 'logits_processor' needs to be on CPU
.OutputMemoryType(OrtMemTypeCPUOutput, 1) // 'sequences_scores' output on CPU
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<MLFloat16>()}),