(device_id=0)
| 34 | |
| 35 | |
| 36 | def get_arch(device_id=0): |
| 37 | compute_cap = 0 |
| 38 | try: |
| 39 | import pynvml |
| 40 | |
| 41 | pynvml.nvmlInit() |
| 42 | handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) |
| 43 | compute_cap = pynvml.nvmlDeviceGetCudaComputeCapability(handle) |
| 44 | compute_cap = compute_cap[0] + compute_cap[1] / 10.0 |
| 45 | except ModuleNotFoundError: |
| 46 | print("NVML not found") |
| 47 | return compute_cap |
| 48 | |
| 49 | |
| 50 | def is_mulit_gpu(): |
no test coverage detected