Recursively move all tensors in a data structure to a device.
(
data: Any,
device: torch.device,
non_blocking: bool = False,
)
| 7 | |
| 8 | |
| 9 | def recursive_to_device( |
| 10 | data: Any, |
| 11 | device: torch.device, |
| 12 | non_blocking: bool = False, |
| 13 | ) -> Any: |
| 14 | """ |
| 15 | Recursively move all tensors in a data structure to a device. |
| 16 | """ |
| 17 | if hasattr(data, "to"): |
| 18 | return data.to(device, non_blocking=non_blocking) |
| 19 | elif isinstance(data, (list, tuple)): |
| 20 | return type(data)(recursive_to_device(d, device, non_blocking) for d in data) |
| 21 | elif isinstance(data, dict): |
| 22 | return {k: recursive_to_device(v, device, non_blocking) for k, v in data.items()} |
| 23 | else: |
| 24 | return data |
| 25 | |
| 26 | |
| 27 | def load_balanced_group_indices( |
no test coverage detected