From 61bcdd177a6317676fc02ee6d306b2ece1b57828 Mon Sep 17 00:00:00 2001
From: Matthew Brookhart <matthewbrookhart@gmail.com>
Date: Tue, 11 Sep 2018 22:14:52 -0700
Subject: [PATCH] initialize visited set with snodes

---
 src/operator/subgraph/partition_graph.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/operator/subgraph/partition_graph.cc b/src/operator/subgraph/partition_graph.cc
index 660ba2aa80a0..cc64fb427a15 100644
--- a/src/operator/subgraph/partition_graph.cc
+++ b/src/operator/subgraph/partition_graph.cc
@@ -239,7 +239,7 @@ bool LabelSubgraph(const Graph& g,
                          const std::vector<nnvm::Node*>& snodes) {
     if (ancestor == descendant) return true;
     std::stack<const nnvm::Node*> s;
-    std::unordered_set<const nnvm::Node*> visited;
+    std::unordered_set<const nnvm::Node*> visited(snodes.begin(), snodes.end());
     s.push(descendant);
     size_t count = 0;
     while (!s.empty()) {
@@ -253,6 +253,8 @@ bool LabelSubgraph(const Graph& g,
         return true;
       }
       for (const auto& entry : top->inputs) {
+        // when searching for the ancestor, the path cannot cross any subgraph node
+        // or a previously visited node
         if (visited.count(entry.node.get()) == 0) {
           s.push(entry.node.get());
         }