Change the device of object recursively
(data, device=None, dtype=None, non_blocking=False, copy=False)
| 6 | |
| 7 | |
| 8 | def to_device(data, device=None, dtype=None, non_blocking=False, copy=False): |
| 9 | """Change the device of object recursively""" |
| 10 | if isinstance(data, dict): |
| 11 | return {k: to_device(v, device, dtype, non_blocking, copy) for k, v in data.items()} |
| 12 | elif dataclasses.is_dataclass(data) and not isinstance(data, type): |
| 13 | return type(data)( |
| 14 | *[to_device(v, device, dtype, non_blocking, copy) for v in dataclasses.astuple(data)] |
| 15 | ) |
| 16 | # maybe namedtuple. I don't know the correct way to judge namedtuple. |
| 17 | elif isinstance(data, tuple) and type(data) is not tuple: |
| 18 | return type(data)(*[to_device(o, device, dtype, non_blocking, copy) for o in data]) |
| 19 | elif isinstance(data, (list, tuple)): |
| 20 | return type(data)(to_device(v, device, dtype, non_blocking, copy) for v in data) |
| 21 | elif isinstance(data, np.ndarray): |
| 22 | return to_device(torch.from_numpy(data), device, dtype, non_blocking, copy) |
| 23 | elif isinstance(data, torch.Tensor): |
| 24 | return data.to(device, dtype, non_blocking, copy) |
| 25 | else: |
| 26 | return data |
| 27 | |
| 28 | |
| 29 | def force_gatherable(data, device): |
no outgoing calls
no test coverage detected
searching dependent graphs…