Get available memory for each GPU.
(max_gpus=None)
| 116 | |
| 117 | |
| 118 | def get_gpu_memory(max_gpus=None): |
| 119 | """Get available memory for each GPU.""" |
| 120 | gpu_memory = [] |
| 121 | num_gpus = ( |
| 122 | torch.cuda.device_count() |
| 123 | if max_gpus is None |
| 124 | else min(max_gpus, torch.cuda.device_count()) |
| 125 | ) |
| 126 | |
| 127 | for gpu_id in range(num_gpus): |
| 128 | with torch.cuda.device(gpu_id): |
| 129 | device = torch.cuda.current_device() |
| 130 | gpu_properties = torch.cuda.get_device_properties(device) |
| 131 | total_memory = gpu_properties.total_memory / (1024 ** 3) |
| 132 | allocated_memory = torch.cuda.memory_allocated() / (1024 ** 3) |
| 133 | available_memory = total_memory - allocated_memory |
| 134 | gpu_memory.append(available_memory) |
| 135 | return gpu_memory |
| 136 | |
| 137 | |
| 138 | def violates_moderation(text): |
nothing calls this directly
no outgoing calls
no test coverage detected