Get available memory for each GPU.
(max_gpus=None)
| 127 | |
| 128 | |
| 129 | def get_gpu_memory(max_gpus=None): |
| 130 | """Get available memory for each GPU.""" |
| 131 | import torch |
| 132 | |
| 133 | gpu_memory = [] |
| 134 | num_gpus = ( |
| 135 | torch.cuda.device_count() |
| 136 | if max_gpus is None |
| 137 | else min(max_gpus, torch.cuda.device_count()) |
| 138 | ) |
| 139 | |
| 140 | for gpu_id in range(num_gpus): |
| 141 | with torch.cuda.device(gpu_id): |
| 142 | device = torch.cuda.current_device() |
| 143 | gpu_properties = torch.cuda.get_device_properties(device) |
| 144 | total_memory = gpu_properties.total_memory / (1024**3) |
| 145 | allocated_memory = torch.cuda.memory_allocated() / (1024**3) |
| 146 | available_memory = total_memory - allocated_memory |
| 147 | gpu_memory.append(available_memory) |
| 148 | return gpu_memory |
| 149 | |
| 150 | |
| 151 | def oai_moderation(text): |
no outgoing calls
no test coverage detected
searching dependent graphs…