torch.device()

torch.device代表将torch.Tensor分配到的设备的对象,有cpu和cuda两种,这里的cuda就是gpu,至于为什么不直接用gpu与cpu对应,是因为gpu的编程接口采用的是cuda

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

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

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

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

device = torch.device('cuda')
#将数据转移到GPU

device = torch.device('cpu')
#将数据转移的cpu

你可能感兴趣的