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)