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)