Pytorch错误记录，关于多GPU训练的问题

刚刚遇到一个极其奇葩的问题。我搭建好了模型，先是在一个GPU下，以如下形式进行了测试

model = xxxNet()
model.train()
model.cuda()

x = torch.randn(1,3,320,640)
model(x)

代码正确运行，说明在单GPU下，模型搭建的没问题。满心欢喜，要开始训练了。然后用了DataParallel，结果报了如下错误。

RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument

错误描述没有截取全。大致说某一个卷积层的输入tensor和weight不是在一个GPU上。
然后我就纳了闷了，数据都是有DataParallel自动去分配到各个GPU上的啊，为啥提示这种错误。

开始在网上找答案，说法不一，都不适合我。
然后终于意识到，DataParallel仅仅支持forward的并行，但我在model的构造函数中是这么写的。

class xxxNet(Module):
	def __init__()
		super.........
		...
		if self.training:
			self.forward = self.training_forward
		else:
			self.forward = self.test_forward

我的本意是想根据是否训练，采用两种不同的前向传播。然后我做了如下修改，成功运行了训练代码

def self.forward(self, x):
	return self.training_forward(x)

这样就跑通了，说明在使用多GPU，一定要让代码通过forward计算前向过程。因为DataParallel仅支持forward的并行，有其他函数想并行，一定要显式的写进forward中。