From 1622cb99372d8c41eede080220315ac165feb870 Mon Sep 17 00:00:00 2001 From: hjchen2 Date: Mon, 19 Nov 2018 13:34:14 +0000 Subject: [PATCH] Fix alpha tensor key --- paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc | 5 ++++- .../fluid/inference/tensorrt/convert/test_leaky_relu_op.cc | 4 ++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc b/paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc index 810295e1919..b3244ef84d8 100644 --- a/paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc @@ -72,7 +72,10 @@ class LeakyReluOpConverter : public OpConverter { nvinfer1::ElementWiseOperation::kSUM); PADDLE_ENFORCE(nullptr != output_layer); // keep alpha tensor to avoid release it's memory - engine_->weight_map[op_desc.Input("alpha")[0]] = std::move(alpha_tensor); + std::string alpha_name = op_desc.Output("Out")[0] + "_alpha"; + PADDLE_ENFORCE(engine_->weight_map.find(alpha_name) == + engine_->weight_map.end()); + engine_->weight_map[alpha_name] = std::move(alpha_tensor); std::string layer_name = "leaky_relu (Output: "; auto output_name = op_desc.Output("Out")[0]; diff --git a/paddle/fluid/inference/tensorrt/convert/test_leaky_relu_op.cc b/paddle/fluid/inference/tensorrt/convert/test_leaky_relu_op.cc index 6fcf78abe43..d00826af075 100644 --- a/paddle/fluid/inference/tensorrt/convert/test_leaky_relu_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/test_leaky_relu_op.cc @@ -20,8 +20,8 @@ namespace paddle { namespace inference { namespace tensorrt { -TEST(leaky_relu_op, test_channel_wise) { - std::unordered_set parameters({"leaky_relu_alpha"}); +TEST(leaky_relu_op, test_leaky_relu) { + std::unordered_set parameters; framework::Scope scope; TRTConvertValidation validator(10, parameters, scope, 1000); validator.DeclInputVar("leaky_relu_input", nvinfer1::DimsCHW(3, 2, 2)); -- GitLab