'cuda:0') does not put layers in lists to GPU

I have the following model where encoder2 is a list of layers (not sequential). However, feeding the data on a GPU through encoder2 results in the data not being in GPU.

class model(nn.Module):
    def __init__(self, args):
        super(model, self).__init__()
        self.encoder1 = make_mlp([6, 128], ["relu"])
        self.encoder2 = [make_mlp([6, 128], ["relu"]) for _ in range(2)]

    def forward(self, data, mode):
        data = torch.stack([self.encoder2[i](data[i])] for i in range(2))       

def make_mlp(dim_list, activations, batch_norm=False, dropout=0):

    if len(dim_list) == 0 and len(activations) == 0:
        return nn.Identity()

    assert len(dim_list) == len(activations)+1
    layers = []
    for dim_in, dim_out, activation in zip(dim_list[:-1], dim_list[1:], activations):
        layers.append(nn.Linear(dim_in, dim_out))
        if batch_norm:
        if activation == 'relu':
        elif activation == 'leakyrelu':
        if dropout > 0:
    return nn.Sequential(*layers)

How can I fix this?

net = model()
net ='cuda:0')
data = data.cuda()
data = net(data)

pytorch has a construct called nn.ModuleList that is exactly for your use case. please have a look here

1 Like