pytorch中sub_module无法转移到gpu问题

通常pytorch在使用net.to('cuda')时可以将net中的__init__函数中的sub_modules, parameters, buffers转到gpu上,但是常规的成员无法转移,例如:

class Model(nn.Module):
	def __init__(self):
    	self.sub_modules = [nn.Conv2d() for i in range(6)]

此时即使调用net.to('cuda'), 成员函数sub_modules中的conv的参数仍然在cpu上,因此会报错:

RuntimeError: Tensor for argument #2 'weight' is on CPU, but expected it to be on GPU.

改正方法为:

class Model(nn.Module):
	def __init__(self):
    	self.device = "cuda" if torch.cuda.is_available() else "cpu"
    	self.sub_modules = [nn.Conv2d().to(device) for i in range(6)]

还有一种错误是将模型的模块没有写在__init__函数中,而是定义在forward函数中,例如:

class Model(nn.Module):
	def __init__(self):
    	pass

    def forward(self, x):
    	out = nn.Conv2D(...)(x)

        return out

也会报上面的错误,而且这种写法也极不规范,正确写法为:

class Model(nn.Module):
	def __init__(self):
    	self.conv = nn.Conv2D(...)

    def forward(self, x):
    	out = self.conv(x)

        return out