Query used GPU memory with pynvml

pynvml_query_memory.py

import datetime
import getopt
import os
import sys
import time

import pynvml

def print_mem(dev_count):
    for index in range(dev_count):
        mem = pynvml.nvmlDeviceGetMemoryInfo(
            pynvml.nvmlDeviceGetHandleByIndex(index)
        )
        used = mem.used // (1024 ** 2)
        total = mem.total // (1024 ** 2)
        print("GPU #%d: %d MB/%d MB" % (index, used, total))

def run_loop(interval):
    pynvml.nvmlInit()
    dev_count = pynvml.nvmlDeviceGetCount()

    while True:
        os.system("clear")
        print(datetime.datetime.now())
        print_mem(dev_count)
        time.sleep(interval)

def main(argv):
    try:
        opts, args = getopt.getopt(argv, "hi:", ["interval="])
    except getopt.GetoptError:
        print("test.py -i <interval>")
        sys.exit(2)
    for opt, arg in opts:
        if opt == "-h":
            print("test.py -i <interval>")
            sys.exit()
        elif opt in ("-i", "--interval"):
            interval = int(arg)
    print(interval)
    run_loop(interval)

if __name__ == "__main__":
    main(sys.argv[1:])