torch.device()基础用法


用途:用于指定将数据或者模型放在GPU/CPU

#1.通常用法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data = data.to(device)
model = model.to(device)
'''
    1.先创建device对象
    2.to函数指定数据或者模型放到哪里
'''

#2.将构建的tensor或者模型放到指定设备上(GPU)
torch.device('cuda',0) #这里的0指的是设备的序号
torch.device('cuda:0')