Pytorch测试GPU是否可用

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

import torch

# 返回当前设备索引
torch.cuda.current_device()

# 返回gpu数量
torch.cuda.device_count()

# 返回gpu名称，索引从0开始
torch.cuda.get_device_name(0)

# cuda是否可用
torch.cuda.is_available()