fix place device (#8450)

starimpact · Oct 28, 2017 · 68038ad · 68038ad
1 parent eb08012
commit 68038ad
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
@@ -639,8 +639,15 @@ inline std::vector<Context> PlaceDevice(const nnvm::IndexedGraph& idx) {
       auto fwd_nid = idx.node_id(idx[i].source->control_deps[0].get());
       CHECK_EQ(idx[fwd_nid].source->op(), _copyto);
       vctx[i] = vctx[idx[fwd_nid].inputs[0].node_id];
-    } else if (idx[i].inputs.size() && vctx[i].dev_type == -1) {
-      vctx[i] = vctx[idx[i].inputs[0].node_id];
+    } else if (idx[i].control_deps.size() &&
+               vctx[idx[i].control_deps[0]].dev_type != -1) {
+      vctx[i] = vctx[idx[i].control_deps[0]];
+    } else {
+      for (const auto& in : idx[i].inputs) {
+        if (vctx[in.node_id].dev_type == -1) continue;
+        vctx[i] = vctx[in.node_id];
+        break;
+      }
     }
   }
   // backward pass