Fix operator initialization order (#15445)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15445

Initilize task graph after operators (task graph uses ops)

Reviewed By: yinghai

Differential Revision: D13530864

fbshipit-source-id: fdc91e9158c1b50fcc96fd1983fd000fdf20c7da

diff --git a/caffe2/core/net_parallel.cc b/caffe2/core/net_parallel.cc
index b9a9f08..3b66ebf 100644 (file)
--- a/caffe2/core/net_parallel.cc
+++ b/caffe2/core/net_parallel.cc
@@ -20,8 +20,6 @@ ParallelNet::ParallelNet(
       num_workers_, 0, "Expected positive number of worker threads");
 
   helper_ = caffe2::make_unique<ParallelNetExecutorHelper>(this);
-  task_graph_ = TaskGraphRegistry()->Create(
-      FLAGS_caffe2_task_graph_engine, helper_.get(), options_);
 
   // initialize operators
   operator_nodes_ = dag_utils::prepareOperatorNodes(net_def, ws);
@@ -32,6 +30,10 @@ ParallelNet::ParallelNet(
     operators_.push_back(op);
   }
 
+  task_graph_ = TaskGraphRegistry()->Create(
+      FLAGS_caffe2_task_graph_engine, helper_.get(), options_);
+  CAFFE_ENFORCE(task_graph_, "Couldn't initialize task graph");
+
   // compute chains
   // TODO: inference mode for chaining
   auto execution_chains = dag_utils::computeChains(operator_nodes_);