测试pytorch（GPU）安装，并输出版本号

不会写代码！！2024-02-01 10:01

GPU版本pytorch安装测试

我们一般完成pytorch(GPU)版本后，要进行测试，以下代码用于备注：

python 复制代码

import torch
flag = torch.cuda.is_available()
print(flag)

ngpu= 1
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.__version__)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())

测试结果如下：

True

cuda:0

2.0.0+cu118

NVIDIA GeForce RTX 3090

tensor([[0.0756, 0.7062, 0.6051],

[0.2772, 0.5758, 0.3481],

[0.4451, 0.9666, 0.9337]], device='cuda:0')

CUDA信息显示

在CUDA后续的更新中，windows系统的任务管理器看不见CUDA占用了，在Ubuntu中同样无法看见，需要在命令行窗口输入：

bash 复制代码

nvidia-smi