(inputs, device_ids)
| 80 | |
| 81 | |
| 82 | def _async_copy_stream(inputs, device_ids): |
| 83 | nr_devs = len(device_ids) |
| 84 | assert type(inputs) in (tuple, list) |
| 85 | assert len(inputs) == nr_devs |
| 86 | |
| 87 | outputs = [] |
| 88 | streams = [_get_stream(d) for d in device_ids] |
| 89 | for i, dev, stream in zip(inputs, device_ids, streams): |
| 90 | with cuda.device(dev): |
| 91 | main_stream = cuda.current_stream() |
| 92 | with cuda.stream(stream): |
| 93 | outputs.append(async_copy_to(i, dev, main_stream=main_stream)) |
| 94 | main_stream.wait_stream(stream) |
| 95 | |
| 96 | return outputs |
| 97 | |
| 98 | |
| 99 | """Adapted from: torch/nn/parallel/_functions.py""" |
no test coverage detected