Returns currently selected device (gpu/cpu). If cuda available, return gpu, otherwise return cpu.
()
| 89 | |
| 90 | |
| 91 | def get_current_device() -> torch.device: |
| 92 | """ |
| 93 | Returns currently selected device (gpu/cpu). |
| 94 | If cuda available, return gpu, otherwise return cpu. |
| 95 | """ |
| 96 | if torch.cuda.is_available(): |
| 97 | return torch.device(f"cuda:{torch.cuda.current_device()}") |
| 98 | else: |
| 99 | return torch.device("cpu") |
| 100 | |
| 101 | |
| 102 | def get_batch_size(data): |
no outgoing calls
no test coverage detected