From 137d6563fc21c472d9256775ecef745c47ea775f Mon Sep 17 00:00:00 2001
From: Zhang Ting <709968123@qq.com>
Date: Thu, 12 Mar 2020 11:55:58 +0800
Subject: [PATCH] add check for assigned data, test=develop (#22960)

---
 paddle/fluid/inference/api/helper.h | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/paddle/fluid/inference/api/helper.h b/paddle/fluid/inference/api/helper.h
index b58c300c2e..cddb0c8daf 100644
--- a/paddle/fluid/inference/api/helper.h
+++ b/paddle/fluid/inference/api/helper.h
@@ -180,11 +180,29 @@ int VecReduceToInt(const std::vector<T> &v) {
   return std::accumulate(v.begin(), v.end(), 1, [](T a, T b) { return a * b; });
 }
 
+template <typename T>
+void CheckAssignedData(const std::vector<std::vector<T>> &data,
+                       const int num_elems) {
+  int num = 0;
+  for (auto it = data.begin(); it != data.end(); ++it) {
+    num += (*it).size();
+  }
+  PADDLE_ENFORCE_EQ(
+      num, num_elems,
+      platform::errors::OutOfRange(
+          "The number of elements out of bounds. "
+          "Expected number of elements = %d. But received %d. Suggested Fix: "
+          "If the tensor is expected to assign %d elements, check the number "
+          "of elements of your 'infer_data'.",
+          num_elems, num, num_elems));
+}
+
 template <typename T>
 static void TensorAssignData(PaddleTensor *tensor,
                              const std::vector<std::vector<T>> &data) {
   // Assign buffer
   int num_elems = VecReduceToInt(tensor->shape);
+  CheckAssignedData(data, num_elems);
   tensor->data.Resize(sizeof(T) * num_elems);
   int c = 0;
   for (const auto &f : data) {
-- 
GitLab