Skip to content

Commit e51d045

Browse files
committed
modify trt engine op test
1 parent e8954a3 commit e51d045

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

paddle/fluid/operators/tensorrt_engine_op_test.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ TEST(TensorRTEngineOp, manual) {
103103
SetAttr<std::string>(engine_op_desc.Proto(), "engine_uniq_key", "a_engine");
104104
SetAttr<std::vector<std::string>>(engine_op_desc.Proto(), "parameters",
105105
std::vector<std::string>({}));
106+
SetAttr<std::vector<std::string>>(engine_op_desc.Proto(), "output_name_mapping",
107+
std::vector<std::string>({"z0"}));
106108

107109
LOG(INFO) << "create engine op";
108110
auto engine_op = framework::OpRegistry::CreateOp(*engine_op_desc.Proto());
@@ -196,6 +198,9 @@ void Execute(int batch_size, int input_dim, int output_dim, int nlayers = 1) {
196198
std::vector<std::string>({"y0", "y1", "y2", "y3"}));
197199
SetAttr<std::string>(engine_op_desc.Proto(), "engine_uniq_key", "b_engine");
198200

201+
SetAttr<std::vector<std::string>>(engine_op_desc.Proto(), "output_name_mapping",
202+
std::vector<std::string>({"z3"}));
203+
199204
auto engine_op = framework::OpRegistry::CreateOp(*engine_op_desc.Proto());
200205

201206
// Execute them.

0 commit comments

Comments
 (0)