Python 利用 Pytorch 获取当前GPU信息

import torch

def get_gpu_info():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        # 获取当前GPU名字
        gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
        # 获取当前GPU总显存
        props = torch.cuda.get_device_properties(device)
        total_memory = props.total_memory / 1e9

        gpu_info = "当前GPU 型号是：{}，可用总显存为：{} GB".format(gpu_name, total_memory)
        return gpu_info, gpu_name
    else:
        src = "No GPU found"
        return src, src

gpu_info, gpu_name = get_gpu_info()

print(gpu_name)