From c1bda4c1cce4d0b71f6d27833b0ad900560e7db2 Mon Sep 17 00:00:00 2001
From: Ye Wang <52801275+wangyems@users.noreply.github.com>
Date: Wed, 23 Nov 2022 10:53:53 -0800
Subject: [PATCH] fix buffer overuse in addtofeed() (#13733)

### Description
<!-- Describe your changes. -->



### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
---
 .../cuda/transformers/generation_device_helper.cc           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
index 11dabc4e97..e5e75a65cc 100644
--- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
+++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
@@ -99,7 +99,7 @@ Status AddToFeeds(const IExecutionProvider* execution_provider,
   size_t total_bytes = 0;
   for (auto& input : inputs) {
     if (input.IsAllocated()) {
-      total_bytes += input.Get<Tensor>().Shape().Size() * input.Type()->Size();
+      total_bytes += input.Get<Tensor>().SizeInBytes();
     }
   }
 
@@ -115,7 +115,7 @@ Status AddToFeeds(const IExecutionProvider* execution_provider,
   for (auto& input : inputs) {
     if (input.IsAllocated()) {
       const Tensor& tensor = input.Get<Tensor>();
-      const size_t bytes = input.Type()->Size() * tensor.Shape().Size();
+      const size_t bytes = tensor.SizeInBytes();
       MLDataType dataType = tensor.DataType();
       if (dataType == DataTypeImpl::GetType<int32_t>()) {
         memcpy(destination, input.Get<Tensor>().Data<int32_t>(), bytes);
@@ -152,7 +152,7 @@ Status AddToFeeds(const IExecutionProvider* execution_provider,
     if (input.IsAllocated()) {
       const Tensor& tensor = input.Get<Tensor>();
       const TensorShape& shape = tensor.Shape();
-      const size_t bytes = input.Type()->Size() * shape.Size();
+      const size_t bytes = tensor.SizeInBytes();
       MLDataType dataType = tensor.DataType();
 
       OrtValue device_input;