From a891032feee9b55f460ea5032fd3e3acca95050f Mon Sep 17 00:00:00 2001
From: Chen Weihang <chenweihang@baidu.com>
Date: Mon, 1 Mar 2021 19:15:13 +0800
Subject: [PATCH]  [Cherry-pick] Fix dtype unmatched in custom op API #31306

[Cherry-pick] Fix dtype unmatched in custom op API

cherry-pick of #31305
---
 paddle/fluid/extension/include/ext_tensor.h  |  4 ++--
 paddle/fluid/extension/src/ext_tensor.cc     |  6 +++---
 paddle/fluid/framework/custom_tensor_test.cc | 10 +++++-----
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/paddle/fluid/extension/include/ext_tensor.h b/paddle/fluid/extension/include/ext_tensor.h
index 77d4ec36e5b..be492a6d553 100644
--- a/paddle/fluid/extension/include/ext_tensor.h
+++ b/paddle/fluid/extension/include/ext_tensor.h
@@ -57,7 +57,7 @@ class PD_DLL_DECL Tensor {
   /// Reshape must be called before calling
   /// mutable_data() or copy_to(const PlaceType& place)
   /// \param shape The shape to set.
-  void reshape(const std::vector<int>& shape);
+  void reshape(const std::vector<int64_t>& shape);
 
   /// \brief Get the memory pointer in CPU or GPU with
   /// specific data type.
@@ -90,7 +90,7 @@ class PD_DLL_DECL Tensor {
   Tensor copy_to(const PlaceType& place) const;
 
   /// \brief Return the shape of the Tensor.
-  std::vector<int> shape() const;
+  std::vector<int64_t> shape() const;
 
   /// \brief Return the data type of the tensor.
   /// It's usually used to get the output tensor data type.
diff --git a/paddle/fluid/extension/src/ext_tensor.cc b/paddle/fluid/extension/src/ext_tensor.cc
index 88c2050bc8a..4434a3bf594 100644
--- a/paddle/fluid/extension/src/ext_tensor.cc
+++ b/paddle/fluid/extension/src/ext_tensor.cc
@@ -95,7 +95,7 @@ void GpuCopy(T *src, T *dst, PlaceType src_plc, PlaceType dst_plc,
   }                                                     \
   auto *tensor = static_cast<framework::LoDTensor *>(tensor_.get());
 
-void Tensor::reshape(const std::vector<int> &shape) {
+void Tensor::reshape(const std::vector<int64_t> &shape) {
   GET_CASTED_TENSOR
   tensor->Resize(framework::make_ddim(shape));
 }
@@ -251,9 +251,9 @@ template PD_DLL_DECL int16_t *Tensor::mutable_data<int16_t>(
     const PlaceType &place);
 template PD_DLL_DECL bool *Tensor::mutable_data<bool>(const PlaceType &place);
 
-std::vector<int> Tensor::shape() const {
+std::vector<int64_t> Tensor::shape() const {
   GET_CASTED_TENSOR
-  return framework::vectorize<int>(tensor->dims());
+  return framework::vectorize<int64_t>(tensor->dims());
 }
 
 const PlaceType &Tensor::place() const {
diff --git a/paddle/fluid/framework/custom_tensor_test.cc b/paddle/fluid/framework/custom_tensor_test.cc
index b891975b96d..2e42248f64b 100644
--- a/paddle/fluid/framework/custom_tensor_test.cc
+++ b/paddle/fluid/framework/custom_tensor_test.cc
@@ -20,7 +20,7 @@
 
 template <typename T>
 paddle::Tensor InitCPUTensorForTest() {
-  std::vector<int> tensor_shape{5, 5};
+  std::vector<int64_t> tensor_shape{5, 5};
   auto t1 = paddle::Tensor(paddle::PlaceType::kCPU);
   t1.reshape(tensor_shape);
   auto* p_data_ptr = t1.mutable_data<T>(paddle::PlaceType::kCPU);
@@ -54,7 +54,7 @@ void TestCopyTensor() {
 }
 
 void TestAPIPlace() {
-  std::vector<int> tensor_shape = {5, 5};
+  std::vector<int64_t> tensor_shape = {5, 5};
 #ifdef PADDLE_WITH_CUDA
   auto t1 = paddle::Tensor(paddle::PlaceType::kGPU);
   t1.reshape(tensor_shape);
@@ -68,7 +68,7 @@ void TestAPIPlace() {
 }
 
 void TestAPISizeAndShape() {
-  std::vector<int> tensor_shape = {5, 5};
+  std::vector<int64_t> tensor_shape = {5, 5};
   auto t1 = paddle::Tensor(paddle::PlaceType::kCPU);
   t1.reshape(tensor_shape);
   CHECK_EQ(t1.size(), 25);
@@ -77,7 +77,7 @@ void TestAPISizeAndShape() {
 
 template <typename T>
 paddle::DataType TestDtype() {
-  std::vector<int> tensor_shape = {5, 5};
+  std::vector<int64_t> tensor_shape = {5, 5};
   auto t1 = paddle::Tensor(paddle::PlaceType::kCPU);
   t1.reshape(tensor_shape);
   t1.template mutable_data<T>();
@@ -86,7 +86,7 @@ paddle::DataType TestDtype() {
 
 template <typename T>
 void TestCast(paddle::DataType data_type) {
-  std::vector<int> tensor_shape = {5, 5};
+  std::vector<int64_t> tensor_shape = {5, 5};
   auto t1 = paddle::Tensor(paddle::PlaceType::kCPU);
   t1.reshape(tensor_shape);
   t1.template mutable_data<T>();
-- 
GitLab