Pytorch——cuda的使用

import torch

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

结果：

True

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

结果：

0

#返回GPU的数量
torch.cuda.device_count()

结果：

1

#返回gpu名字，设备索引默认从0开始
torch.cuda.get_device_name(0)

'NVIDIA TITAN Xp'

#设置当前设备。
device = "cuda:0"
torch.cuda.set_device(device)