指定GPU的编号
网络,损失函数,训练的数据,标签需要调入GPU中训练
to.(device)
import torch
"""
net:网络
Loss:loss损失函数
data:数据
label:标签
"""
device = torch.device("cuda:0")
print(device)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
data = data.to(device)