PyTorch获取GPU信息（设备id、名称、显存）

在PyTorch中，我们可以使用如下代码获取GPU信息：

import torch


def gpu_info() -> str:
    info = ''
    for id in range(torch.cuda.device_count()):
        p = torch.cuda.get_device_properties(id)
        info += f'CUDA:{id} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n'
    return info[:-1]


if __name__ == '__main__':
    print(gpu_info())