Send each torch.Tensor in arr to device.
(
arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]],
device: torch.device,
)
| 74 | |
| 75 | |
| 76 | def to_device( |
| 77 | arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]], |
| 78 | device: torch.device, |
| 79 | ) -> T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]]: |
| 80 | """ |
| 81 | Send each torch.Tensor in arr to device. |
| 82 | """ |
| 83 | if isinstance(arr, torch.Tensor): |
| 84 | arr = arr.to(device=device) |
| 85 | return arr |
| 86 | elif isinstance(arr, (list, tuple)): |
| 87 | return [to_device(x, device=device) for x in arr] |
| 88 | elif isinstance(arr, dict): |
| 89 | for key, val in arr.items(): |
| 90 | arr[key] = to_device(val, device=device) |
| 91 | return arr |
| 92 | else: |
| 93 | return arr |
| 94 | |
| 95 | |
| 96 | def to_dtype( |