Return true iff (dst, dst_input) is specified on host memory.
| 134 | |
| 135 | // Return true iff (dst, dst_input) is specified on host memory. |
| 136 | bool IsDstInputOnHost(const Edge* edge, const GraphInfo& info) { |
| 137 | const Node* dst = edge->dst(); |
| 138 | int dst_port = edge->dst_input(); |
| 139 | if (info.device_types[dst->id()] != DEVICE_CPU) { |
| 140 | if (edge->IsControlEdge()) return false; |
| 141 | auto dst_it = info.input_types.find({dst->id(), dst_port}); |
| 142 | DCHECK(dst_it != info.input_types.end()); |
| 143 | return dst_it->second == HOST_MEMORY; |
| 144 | } |
| 145 | return true; |
| 146 | } |
| 147 | |
| 148 | void SetInput(NodeDef* dst, int dst_slot, |
| 149 | StringPiece src_name, int src_slot) { |