From 5f4e0c95ec1ea21baa788102f672a949e1b3f8b4 Mon Sep 17 00:00:00 2001
From: Hariharan Seshadri <shariharan91@gmail.com>
Date: Sun, 4 Dec 2022 21:30:57 -0800
Subject: [PATCH] Misc minor bug fixes in transformer kernels (#13780)

---
 .../contrib_ops/cpu/transformers/beam_search_impl_base.h      | 2 +-
 onnxruntime/contrib_ops/cuda/bert/attention.cc                | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
index 174d78af4f..f3764fcbfa 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
@@ -191,7 +191,7 @@ Status BeamSearchBase<T>::CheckInputs(const OpKernelContextInternal& context) {
                                             context.Input<Tensor>(0),     // input_ids
                                             context.Input<Tensor>(7),     // vocab_mask
                                             context.Input<Tensor>(8),     // prefix_vocab_mask
-                                            context.Input<Tensor>(10)));  // attention_mask
+                                            context.Input<Tensor>(9)));  // attention_mask
 
   return Status::OK();
 }
diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc
index 158b3f7c75..cd13ed8bdb 100644
--- a/onnxruntime/contrib_ops/cuda/bert/attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc
@@ -122,14 +122,14 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
   cublasHandle_t cublas = CublasHandle();
   constexpr size_t element_size = sizeof(T);
 
-  IAllocatorUniquePtr<T> gemm_buffer;
+  IAllocatorUniquePtr<void> gemm_buffer;
   if (weights != nullptr) {
     // Use GEMM for fully connection.
     int m = batch_size * sequence_length;
     int n = (parameters.hidden_size + parameters.hidden_size + parameters.v_hidden_size);
     int k = parameters.input_hidden_size;
     size_t gemm_buffer_size = static_cast<size_t>(batch_size) * sequence_length * n * element_size;
-    gemm_buffer = GetScratchBuffer<T>(gemm_buffer_size);
+    gemm_buffer = GetScratchBuffer<void>(gemm_buffer_size);
 
     typedef typename ToCudaType<T>::MappedType CudaT;
     CudaT one = ToCudaType<T>::FromFloat(1.0f);