Get available memory for each GPU.
(max_gpus=None)
| 15 | return system_message |
| 16 | |
| 17 | def get_gpu_memory(max_gpus=None): |
| 18 | """Get available memory for each GPU.""" |
| 19 | gpu_memory = [] |
| 20 | num_gpus = ( |
| 21 | torch.cuda.device_count() |
| 22 | if max_gpus is None |
| 23 | else min(max_gpus, torch.cuda.device_count()) |
| 24 | ) |
| 25 | |
| 26 | for gpu_id in range(num_gpus): |
| 27 | with torch.cuda.device(gpu_id): |
| 28 | device = torch.cuda.current_device() |
| 29 | gpu_properties = torch.cuda.get_device_properties(device) |
| 30 | total_memory = gpu_properties.total_memory / (1024**3) |
| 31 | allocated_memory = torch.cuda.memory_allocated() / (1024**3) |
| 32 | available_memory = total_memory - allocated_memory |
| 33 | gpu_memory.append(available_memory) |
| 34 | return gpu_memory |
| 35 | |
| 36 | |
| 37 | def standardize_category(category): |