(self)
| 55 | self.pynvml = load_pynvml_env() |
| 56 | |
| 57 | def get_cuda_id(self): |
| 58 | if self.cuda_available: |
| 59 | gpu_mem = sorted(self.gpu_mem_get_all(), key=lambda item: item.free, reverse=True) |
| 60 | low_use_id = gpu_mem[0].id |
| 61 | return th.device(f'cuda:{low_use_id}') |
| 62 | else: |
| 63 | return th.device('cpu') |
| 64 | |
| 65 | def gpu_mem_get_all(self): |
| 66 | "get total, used and free memory (in MBs) for each available gpu" |
nothing calls this directly
no test coverage detected