modifying pre/post-partition declaration

apache · Aug 16, 2019 · df67ed0 · df67ed0
1 parent 6667bce
commit df67ed0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
@@ -1200,11 +1200,11 @@ int MXOptimizeForBackend(SymbolHandle sym_handle,
   for (auto property : subgraph_prop_list) {
     nnvm::Graph g = orig_g;
     property->SetAttr("graph", g);
-    property->PrePartition();
+    property->PrePartition(g);
     g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(property);
     g = ApplyPass(std::move(g), "BuildSubgraph");
     g.attrs.erase("subgraph_property");
-    property->PostPartition();
+    property->PostPartition(g);
     property->RemoveAttr("graph");
     s->outputs = g.outputs;
   }

diff --git a/src/operator/subgraph/subgraph_property.h b/src/operator/subgraph/subgraph_property.h
@@ -221,9 +221,9 @@ class SubgraphProperty {
     return nullptr;
   }
 
-  virtual PrePartition() {}
+  virtual void PrePartition(const nnvm::Graph g) {}
 
-  virtual PostPartition() {}
+  virtual void PostPartition(const nnvm::Graph g) {}
 
   virtual SubgraphSelectorV2Ptr CreateSubgraphSelectorV2() const {
     auto v1_ptr = CreateSubgraphSelector();