Sets the device field of ops in the graph.
| 21 | |
| 22 | // Sets the device field of ops in the graph. |
| 23 | Status SetDevice(const GraphDef& input_graph_def, |
| 24 | const TransformFuncContext& context, |
| 25 | GraphDef* output_graph_def) { |
| 26 | string new_device; |
| 27 | TF_RETURN_IF_ERROR(context.GetOneStringParameter("device", "", &new_device)); |
| 28 | bool if_default; |
| 29 | TF_RETURN_IF_ERROR( |
| 30 | context.GetOneBoolParameter("if_default", false, &if_default)); |
| 31 | |
| 32 | output_graph_def->Clear(); |
| 33 | for (const NodeDef& node : input_graph_def.node()) { |
| 34 | NodeDef* new_node = output_graph_def->mutable_node()->Add(); |
| 35 | *new_node = node; |
| 36 | if (!if_default || (node.device().empty())) { |
| 37 | new_node->set_device(new_device); |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | return Status::OK(); |
| 42 | } |
| 43 | |
| 44 | REGISTER_GRAPH_TRANSFORM("set_device", SetDevice); |
| 45 |