首先将数据转移至cuda
实例化net
net = model_name()
net.cuda()
并行处理数据
net = nn.DataParallel(net)
将损失函数转移至cuda
criterion = nn.MSELoss()
criterion = criterion.cuda()
在训练时,也要将data转移至cuda
data = data.cuda()
label = label.cuda()
首先将数据转移至cuda
实例化net
net = model_name()
net.cuda()
并行处理数据
net = nn.DataParallel(net)
将损失函数转移至cuda
criterion = nn.MSELoss()
criterion = criterion.cuda()
在训练时,也要将data转移至cuda
data = data.cuda()
label = label.cuda()