mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-08-24 11:03:19 +00:00
23 lines
732 B
Python
23 lines
732 B
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
|
|
import torch
|
|
|
|
def get_gpu_memory(max_gpus=None):
|
|
gpu_memory = []
|
|
num_gpus = (
|
|
torch.cuda.device_count()
|
|
if max_gpus is None
|
|
else min(max_gpus, torch.cuda.device_count())
|
|
)
|
|
|
|
for gpu_id in range(num_gpus):
|
|
with torch.cuda.device(gpu_id):
|
|
device = torch.cuda.current_device()
|
|
gpu_properties = torch.cuda.get_device_properties(device)
|
|
total_memory = gpu_properties.total_memory / (1024 ** 3)
|
|
allocated_memory = torch.cuda.memory_allocated() / (1024 ** 3)
|
|
available_memory = total_memory - allocated_memory
|
|
gpu_memory.append(available_memory)
|
|
return gpu_memory
|