Creating a Feed Forward NN Model in Pytorch with a dynamic number of hidden layers
Why are these two code segments not equivalent:
Segment 1: Creating a model with 2 layers.

class FNNModule(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim1, hidden_dim2, non_linear_function):
        super().__init__()
        self.hidden1 = nn.Linear(input_dim, hidden_dim1)
        self.hidden2 = nn.Linear(hidden_dim1, hidden_dim2)       
        self.non_linear_function = non_linear_function()
        self.final_linear = nn.Linear(hidden_dim2, output_dim)

    def forward(self, x):
        out = self.hidden1(x)
        out = self.non_linear_function(out)
        out = self.hidden2(x)
        out = self.non_linear_function(out)        
        out = self.final_linear(out)
        return out


Segment Two: Creating the same model but changing code where hidden_layers is a variable:

class FNNModuleVar(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim_array = [], non_linear_function_array=[]):
        super().__init__()
        self.linear_functions = []
        self.non_linear_functions = [x() for x in non_linear_function_array]
        self.hidden_layers = len(hidden_dim_array)
        for l in range(self.hidden_layers):
            self.linear_functions.append(nn.Linear(input_dim, hidden_dim_array[l]))
            input_dim = hidden_dim_array[l]
        self.final_linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = x
        for i in range(self.hidden_layers):
            out = self.linear_functions[i](out)
            out = self.non_linear_functions[i](out)
        out = self.final_linear(x)
        return out
modelVar = FNNModuleVar(input_dim, output_dim, [100, 50], [nn.Tanh, nn.Tanh])
model = FNNModule(input_dim, output_dim, 100, 50, nn.Tanh)


When I try to iterate through modelVar.parameters() and model.parameters() I see that I have very different models.

What am I doing wrong in modelVar?

Those modules are called as you would expect them to be called they are just not visible to the Module. In order to make them visible you can wrap them in a nn.ModuleList like this:

class FNNModuleVar(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim_array = [], non_linear_function_array=[]):
        super().__init__()
        self.linear_functions = []
        self.non_linear_functions = [x() for x in non_linear_function_array]
        self.hidden_layers = len(hidden_dim_array)
        for l in range(self.hidden_layers):
            self.linear_functions.append(nn.Linear(input_dim, hidden_dim_array[l]))
            input_dim = hidden_dim_array[l]
        self.linear_functions = nn.ModuleList(self.linear_functions)
        self.final_linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = x
        for i in range(self.hidden_layers):
            out = self.linear_functions[i](out)
            out = self.non_linear_functions[i](out)
        out = self.final_linear(out)
        return out


printing the models now would yield:

FNNModule(
  (hidden1): Linear(in_features=50, out_features=100, bias=True)
  (hidden2): Linear(in_features=100, out_features=50, bias=True)
  (non_linear_function): Tanh()
  (final_linear): Linear(in_features=50, out_features=100, bias=True)
)
FNNModuleVar(
  (linear_functions): ModuleList(
    (0): Linear(in_features=50, out_features=100, bias=True)
    (1): Linear(in_features=100, out_features=50, bias=True)
  )
  (final_linear): Linear(in_features=50, out_features=100, bias=True)
)


More details: https://pytorch.org/docs/stable/nn.html#torch.nn.ModuleList
