torch.device使用方法

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