diff --git a/paddle/fluid/inference/tensorrt/convert/fill_constant_batch_size_like_op.cc b/paddle/fluid/inference/tensorrt/convert/fill_constant_batch_size_like_op.cc index 5f00777a663af96ce6b74ff9c11ca82233b8a6e7..b163cd5fec27895b2f9e2a364e2c419130f0ad32 100644 --- a/paddle/fluid/inference/tensorrt/convert/fill_constant_batch_size_like_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/fill_constant_batch_size_like_op.cc @@ -44,6 +44,10 @@ class FillConstantBatchSizeLikeOpConverter : public OpConverter { PADDLE_GET_CONST(std::string, op_desc.GetAttr("str_value")); std::vector shape = PADDLE_GET_CONST(std::vector, op_desc.GetAttr("shape")); + if (str_value == "") { + float value = PADDLE_GET_CONST(float, op_desc.GetAttr("value")); + str_value = std::to_string(value); + } float value = std::stof(str_value); auto* input_shape_tensor = Shape(input); @@ -65,7 +69,7 @@ class FillConstantBatchSizeLikeOpConverter : public OpConverter { auto layer = TRT_ENGINE_ADD_LAYER( engine_, Fill, nvinfer1::Dims{}, nvinfer1::FillOperation::kLINSPACE); std::vector value_vec(1, value); - std::vector beta_vec(3, 0.); + std::vector beta_vec(shape.size(), 0.); layer->setAlpha(value); layer->setBeta(0.f); layer->setInput(0, *out_shape_tensor); diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_rnn.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_rnn.py index 2a3c25bab11ca26795bfac68d86137153bb0f0cb..a6000b1256f93164d6aaffc6a631b490ef9752c5 100644 --- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_rnn.py +++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_rnn.py @@ -87,7 +87,9 @@ class TrtConvertSliceTest(TrtLayerAutoScanTest): "input_dim_idx": 0, "str_value": - "0.0", + "", + "value": + 0.0, "shape": [K * num_layers, -1, hidden_size], "output_dim_idx": 1,