Move inputs to a device
(inputs: Any, device: torch.device)
| 172 | |
| 173 | |
| 174 | def to_device(inputs: Any, device: torch.device) -> Any: |
| 175 | """Move inputs to a device""" |
| 176 | if callable(getattr(inputs, 'to', None)): |
| 177 | return inputs.to(device=device) |
| 178 | |
| 179 | if isinstance(inputs, Mapping): |
| 180 | res = {} |
| 181 | for k, v in inputs.items(): |
| 182 | res[k] = to_device(v, device) |
| 183 | elif isinstance(inputs, Sequence) and not isinstance(inputs, str): |
| 184 | res = [] |
| 185 | for b in inputs: |
| 186 | res.append(to_device(b, device)) |
| 187 | else: |
| 188 | res = inputs |
| 189 | return res |
| 190 | |
| 191 | |
| 192 | def upper_bound(lo: int, hi: int, cond: Callable[[int], bool]) -> int: |