Model.to(device) does't work in dynamic graph?

I create a dynamic conv-net by for in range, and data.to(device) as well as model.to(device).
However it raise ‘RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same’.
Is there something wrong when i create dynamic net and use cuda?
Here is the key code, thanks very much!

# create dynamic net
class AutoNet(nn.Module):
    def __init__(self, channels=[32, 32, 32], conv_kernel_size=3, conv_padding=1, pool_kernel_size=2, pool_stride=2):
        super(AutoNet, self).__init__()
        self.pool = nn.MaxPool2d(kernel_size=pool_kernel_size, stride=pool_stride)

        self.convs = []
        channel_in = 1
        layer_out = args.image_size
        for i in range(len(channels)):
            channel = channels[i]
            layer_out = (layer_out - conv_kernel_size + 2 * conv_padding) / 1 + 1
            layer_out = (layer_out - pool_kernel_size) / pool_stride + 1
            self.convs.append(
                nn.Conv2d(in_channels=channel_in, out_channels=channel, kernel_size=conv_kernel_size, padding=conv_padding))
            channel_in = channel
        self.fc = nn.Linear(layer_out * layer_out * channels[-1], args.num_class)

    def forward(self, x):
        for i in range(len(self.convs)):
            x = self.pool(self.convs[i](x))
        x = x.view(args.batch_size, -1)
        x = self.fc(x)
        return x

# model
model = AutoNet(channels=args.channels,
                    conv_kernel_size=args.conv_kernel_size,
                    conv_padding=args.conv_padding,
                    pool_kernel_size=args.pool_kernel_size,
                    pool_stride=args.pool_stride)
device = torch.device('cuda:' + str(args.gpu) if args.gpu != -1 else 'cpu')
model.to(device)

# train
inputs, labels = data[0].to(device), data[1].to(device)
outs = model(inputs)

Hi,

The problem is that you store you convolutions in a plain python list. This means that .to() cannot find these convs to move them to cuda.
You should do self.convs = nn.ModuleList() instead of self.convs = [] and that will work fine.

1 Like