Make inlined functions inherit device.

PiperOrigin-RevId: 158567875
TeeeJay · Jun 9, 2017 · d4ae391 · d4ae391
1 parent f80fd0f
commit d4ae391
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
@@ -850,6 +850,7 @@ void InlineFunctionBody(const FunctionLibraryDefinition& flib_def, Graph* g,
  for (Node* n : fbody->graph->op_nodes()) {
  NodeDef ndef = n->def();
  ndef.set_name(strings::StrCat(caller->name(), "/", ndef.name()));
+ ndef.set_device(caller->def().device());
  Node* clone = g->AddNode(ndef, &s);
  TF_CHECK_OK(s);
  node_map[n->id()] = clone;