From fc96463b25c1f0bf9d48541bdfb2d0f0cf3e082b Mon Sep 17 00:00:00 2001
From: Yang Yang <yangyang62@baidu.com>
Date: Thu, 12 Oct 2017 01:15:37 +0000
Subject: [PATCH] pass multiple target

---
 paddle/framework/prune_test.cc | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/paddle/framework/prune_test.cc b/paddle/framework/prune_test.cc
index ab08b851d3d..790fa169244 100644
--- a/paddle/framework/prune_test.cc
+++ b/paddle/framework/prune_test.cc
@@ -161,3 +161,21 @@ TEST(Prune, multi_output_op) {
   Prune(*pdesc, pruned, 0);
   PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 2);
 }
+
+TEST(Prune, multi_target) {
+  f::ProgramDesc *program_desc = GetNewProgramDesc();
+  f::ProgramDescBind &program = f::ProgramDescBind::Instance(program_desc);
+  f::BlockDescBind *block = program.Block(0);
+
+  AddOp("one_two", {{"input", {"a"}}}, {{"output", {"b", "c"}}}, {}, block);
+  AddOp("one_one", {{"input", {"b"}}}, {{"output", {"b1"}}}, {}, block);
+  AddOp("one_one", {{"input", {"c"}}}, {{"output", {"c1"}}}, {}, block);
+
+  f::ProgramDesc *pdesc = program.Proto();
+  pdesc->mutable_blocks(0)->mutable_ops(1)->set_is_target(true);
+  pdesc->mutable_blocks(0)->mutable_ops(2)->set_is_target(true);
+
+  f::ProgramDesc pruned;
+  Prune(*pdesc, pruned, 0);
+  PADDLE_ENFORCE_EQ(pruned.blocks(0).ops_size(), 3);
+}
-- 
GitLab