(obj, dev, main_stream=None)
| 11 | |
| 12 | |
| 13 | def async_copy_to(obj, dev, main_stream=None): |
| 14 | if torch.is_tensor(obj): |
| 15 | v = obj.cuda(dev, non_blocking=True) |
| 16 | if main_stream is not None: |
| 17 | v.data.record_stream(main_stream) |
| 18 | return v |
| 19 | elif isinstance(obj, collections.Mapping): |
| 20 | return {k: async_copy_to(o, dev, main_stream) for k, o in obj.items()} |
| 21 | elif isinstance(obj, collections.Sequence): |
| 22 | return [async_copy_to(o, dev, main_stream) for o in obj] |
| 23 | else: |
| 24 | return obj |
| 25 | |
| 26 | |
| 27 | def dict_gather(outputs, target_device, dim=0): |
no outgoing calls
no test coverage detected