From 63c9973b7a8194323b00d7d2e25ba180bb9cb3be Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Wed, 24 May 2023 16:13:11 -0700 Subject: [PATCH] Fix cuda provider crash on it (#16056) --- onnxruntime/contrib_ops/cuda/transformers/beam_search.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc b/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc index fe775e1ba2..759102b099 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search.cc @@ -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(), DataTypeImpl::GetTensorType()}),