File size: 598 Bytes
24bc58a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import pynvml

def get_free_vram_gpu0():
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)  # Select GPU 0
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    free_vram_bytes = mem_info.free

    # Convert bytes to human-readable format
    def sizeof_fmt(num, suffix='B'):
        for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
            if abs(num) < 1024.0:
                return "%3.1f %s%s" % (num, unit, suffix)
            num /= 1024.0
        return "%.1f %s%s" % (num, 'Yi', suffix)

    return sizeof_fmt(free_vram_bytes)

print(get_free_vram_gpu0())