diff --git a/paddle/fluid/operators/tensorrt_engine_op_test.cc b/paddle/fluid/operators/tensorrt_engine_op_test.cc index b6e285dc5dc4ae903213718348e7a33f60fae8a7..37657fa0b0498986fe67027415279af1775e58b9 100644 --- a/paddle/fluid/operators/tensorrt_engine_op_test.cc +++ b/paddle/fluid/operators/tensorrt_engine_op_test.cc @@ -103,7 +103,8 @@ TEST(TensorRTEngineOp, manual) { SetAttr(engine_op_desc.Proto(), "engine_uniq_key", "a_engine"); SetAttr>(engine_op_desc.Proto(), "parameters", std::vector({})); - SetAttr>(engine_op_desc.Proto(), "output_name_mapping", + SetAttr>(engine_op_desc.Proto(), + "output_name_mapping", std::vector({"z0"})); LOG(INFO) << "create engine op"; @@ -198,7 +199,8 @@ void Execute(int batch_size, int input_dim, int output_dim, int nlayers = 1) { std::vector({"y0", "y1", "y2", "y3"})); SetAttr(engine_op_desc.Proto(), "engine_uniq_key", "b_engine"); - SetAttr>(engine_op_desc.Proto(), "output_name_mapping", + SetAttr>(engine_op_desc.Proto(), + "output_name_mapping", std::vector({"z3"})); auto engine_op = framework::OpRegistry::CreateOp(*engine_op_desc.Proto());