Pytorch学习笔记(2)——to(device)无法将数据放到GPU上

1 问题来源

我删除了无关代码，只放出错的部分

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_iter = Data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)

model = RNNSa(300, 2).to(device)

print('\nmodel.parameters().device:', next(model.parameters()).device)

for epoch in range(num_epochs):
	for X, y in train_iter:
		X = X.to(device)
		y = y.to(device)
		print('train X device:', X.device)
        print('train y device:', y.device)

整个代码就是我将数据打包到DataLoader()中；接着声明了一个模型的实例化对象，并放到GPU上；然后从DataLoader()中提取X与y，并将这两个放到GPU上。逻辑上没有问题，但是却报错：

RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cpu and parameter tensor at cuda:0

打印信息为：

next(model.parameters()).device: cuda:0
train X device: cpu
train y device: cpu

具体为什么出错我也不是特别清楚，但是由于我们服务器上有两张卡，所以极有可能是model先占用了cuda:0这张卡后，X与y找不到放哪个卡上导致没有放到GPU上。

2 解决方案

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_iter = Data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)

model = RNNSa(300, 2).to(device)
cuda = next(model.parameters()).device

print('\nmodel.parameters().device:', next(model.parameters()).device)

for epoch in range(num_epochs):
	for X, y in train_iter:
		X = X.to(cuda)
        y = y.to(cuda)
		print('train X device:', X.device)
        print('train y device:', y.device)

先将model丢上GPU，接着获取存储model的device，然后将X与y放到这个device上即可解决：

model.parameters().device: cuda:0
train X device: cuda:0
train y device: cuda:0