Does nn.Module not support python list?

1

I have a model class like this, and I tried to run it on gpu:

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(3, 64, kernel_size = 3)
        self.avgpool = nn.AvgPool2d(3, 1)
        self.denses = []
        self.denses.append(nn.Linear(64, 10))
        self.denses.append(nn.Linear(64, 20))

    def forward(self, x):
        x = self.conv(x)
        x = self.avgpool(x).view(-1, 64)
        out1 = self.denses[0](x)
        out2 = self.denses[1](x)
        return out1, out2


m = Model()
in_tensor = torch.randn(1, 3, 128, 128).cuda()
m.cuda()
out = m(in_tensor)

Then I have the error message:

Traceback (most recent call last):
  File "try.py", line 56, in <module>
    out = m(in_tensor)
  File "/home/zhangzy/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "try.py", line 48, in forward
    out1 = self.denses[0](x)
  File "/home/zhangzy/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/zhangzy/.local/lib/python3.5/site-packages/torch/nn/modules/linear.py", line 55, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/zhangzy/.local/lib/python3.5/site-packages/torch/nn/functional.py", line 992, in linear
    return torch.addmm(bias, input, weight.t())
RuntimeError: Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument #4 'mat1'

Can’t I use python list to hold these dense branches? How shall I use structure such as list or tuple to manipulate these layers properly ?

2

Use nn.ModuleList instead of a plain Python list.

3

You might want to look into nn.Sequential() which purpose is to add layers together as you are trying with the list.

Here is an example of a convolutional layer, followed by activation and then a pooling using sequential. Hope it makes sense.

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Sequential(         
            nn.Conv1d(
                in_channels=8,              
                out_channels=16,             
                kernel_size=5,              
                stride=1,                   
                padding=2,                  
            ),                              
            nn.ReLU(),                      
            nn.MaxPool1d(kernel_size=2), 
        ) 
       
    def forward(self, x):
        x = self.conv1(x.double()) #inputs (1,3,batch_size)