add int32 support

PaddlePaddle · Jan 21, 2022 · 132a0eb · 132a0eb
1 parent fec3c5b
commit 132a0eb
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/paddle/fluid/inference/tensorrt/convert/reduce_op.cc b/paddle/fluid/inference/tensorrt/convert/reduce_op.cc
@@ -83,6 +83,7 @@ class ReduceOpConverter : public OpConverter {
     }
 
     auto output_name = op_desc.Output("Out")[0];
+    layer->getOutput(0)->setType(layer->getInput(0)->getType());
     RreplenishLayerAndOutput(layer, op_type, {output_name}, test_mode);
   }
 

diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_reduce_mean.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_reduce_mean.py
@@ -50,7 +50,7 @@ def generate_input1(dtype, attrs: List[Dict[str, Any]]):
             for dim in [[], [1], [0], [0, 1], [1, 2, 3], [-2, 0, 3], [-3],
                         [-4, 1], [3, 4, 5]]:
                 for reduce_all in [True, False]:
-                    for out_dtype in [-1, 5]:
+                    for out_dtype in [-1, 2, 5]:
                         dics = [{
                             "keep_dim": keep_dim,
                             "dim": dim,

diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_reduce_sum.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_reduce_sum.py
@@ -51,7 +51,7 @@ def generate_input1(dtype, attrs: List[Dict[str, Any]]):
             for dim in [[], [1], [0], [0, 1], [1, 2, 3], [-2, 0, 3], [-3],
                         [-4, 1], [3, 4, 5]]:
                 for reduce_all in [True, False]:
-                    for out_dtype in [-1, 5]:
+                    for out_dtype in [-1, 2, 5]:
                         dics = [{
                             "keep_dim": keep_dim,
                             "dim": dim,