Hi I am very new to Pytorch!
I am trying to create a model that allows the user to specify the number of hidden layers to be integrated to the network.
Specifically, this is my model :
class MLP(nn.Module):

    def __init__(self, h_sizes, out_size):

        super(MLP, self).__init__()

        # Hidden layers
        self.hidden = []
        for k in range(len(h_sizes)-1):
            self.hidden.append(nn.Linear(h_sizes[k], h_sizes[k+1]))

        # Output layer
        self.out = nn.Linear(h_sizes[-1], out_size)

    def forward(self, x):

        # Feedforward
        for layer in self.hidden:
            x = F.relu(layer(x))
        output= F.softmax(self.out(x), dim=1)

        return output

I would like to be able to give it the parameter “h_sizes” which would be a list of integers defining the size of each hidden layer (and the length of the list defining the number of hidden layers to build).
However, it doesn’t quite work as it doesn’t seem to include the parameters of those layers in “model.parameters”, making it unconvenient to send the model on GPU.
What’s the right way of doing this?