(proto)
| 1600 | mapping = {} |
| 1601 | |
| 1602 | def map_ops(proto): |
| 1603 | for op in proto.op: |
| 1604 | device_option = op.device_option |
| 1605 | if op.type == "Iter": |
| 1606 | # Hack for Iters which have blob in CPU context |
| 1607 | device_option = caffe2_pb2.DeviceOption() |
| 1608 | device_option.device_type = caffe2_pb2.CPU |
| 1609 | for b in list(op.input) + list(op.output): |
| 1610 | if b not in mapping: |
| 1611 | mapping[b] = device_option |
| 1612 | if op.type.startswith('RecurrentNetwork'): |
| 1613 | step_args = [a for a in op.arg if a.name.endswith("step_net")] |
| 1614 | for step_arg in step_args: |
| 1615 | map_ops(step_arg.n) |
| 1616 | map_ops(model.param_init_net.Proto()) |
| 1617 | map_ops(model.net.Proto()) |
| 1618 | model._blob_to_device = mapping |
no test coverage detected
searching dependent graphs…