How to convert layer list in nn.Module to CPU or GPU?

I build a nn.Module that has a list containing some Linear.
I try to convert it to cuda but got error: RuntimeError: Expected object of backend CPU but got backend CUDA for argument #4 'mat1'

Is there any way to convert it? Here is some example code

class DeepLinear(nn.Module):
    def __init__(self, dim: int, depth: int):
        super(DeepLinear, self).__init__()
        self.depth = depth
        self.linear = []
        for _ in range(depth):
            self.linear.append(nn.Linear(dim, dim))

    def forward(self, inputs):
        for i in range(self.depth):
            inputs = self.linear[i](inputs)

        return inputs

a = DeepLinear(3, 5).cuda()
x = torch.ones(5, 3).cuda()
print(a(x)) # error
1 Like

You are using python list. If you wrap your layers into torch.nn.ModuleList then you can send the model to cuda:

class DeepLinear(nn.Module):
    def __init__(self, dim: int, depth: int):
        super(DeepLinear, self).__init__()
        self.depth = depth
        self.linear = []
        for _ in range(depth):
            self.linear.append(nn.Linear(dim, dim))

        self.linear = nn.ModuleList(self.linear)

    def forward(self, inputs):
        for i in range(self.depth):
            inputs = self.linear[i](inputs)

        return inputs

device = torch.device("cuda:0")
model = DeepLinear(3, 5).to(device)
x = torch.ones(5, 3).to(device)

print(model(x))

which works and prints the following:

tensor([[-0.4190,  0.2145, -0.1509],
        [-0.4190,  0.2145, -0.1509],
        [-0.4190,  0.2145, -0.1509],
        [-0.4190,  0.2145, -0.1509],
        [-0.4190,  0.2145, -0.1509]], device='cuda:0', grad_fn=<AddmmBackward>)
5 Likes