fix a bug

4f71a3b1 · nhzlx · c8adfb34 · 4f71a3b1
隐藏空白更改
内联并排

Showing with 2 addition and 1 deletion

paddle/fluid/inference/tensorrt/convert/ut_helper.h paddle/fluid/inference/tensorrt/convert/ut_helper.h +2 -1

未找到文件。
--- a/paddle/fluid/inference/tensorrt/convert/ut_helper.h
+++ b/paddle/fluid/inference/tensorrt/convert/ut_helper.h
@@ -162,7 +162,8 @@ class TRTConvertValidation {

      size_t fluid_out_size = fluid_out.size();
      if (if_add_batch_ == true) {
-        fluid_out_size = batch_size * (tensor->dims().size() / max_batch_size_);
+        fluid_out_size =
+            batch_size * (framework::product(tensor->dims()) / max_batch_size_);
      }
      // Compare two output
      ASSERT_FALSE(fluid_out.empty());