set variable network layers based on parameters in pytorch
I want to make the following network definition to a parametric one. The number of continuous and discrete columns varies for different data. I first pass the whole input data, which in this case is 110 dimensional , from a linear with a relu activation. The output of each categorical field of my data varies based on a previous one-hot encoding data transformation. I need to define a nn.Linear(110, number of encodings) for each of them.
class Generator(nn.Module):
  def __init__(self):
    super(Generator, self).__init__(110)
    self.lin1 = nn.Linear(110,110)
    self.lin_numerical = nn.Linear(110, 6)
    self.lin_cat_job = nn.Linear(110, 9)
    self.lin_cat_sex = nn.Linear(110, 2)
    self.lin_cat_incomeclass = nn.Linear(110, 7)

  def forward(self, x):
    x = torch.relu(self.lin1(x))
    x_numerical = f.leaky_relu(self.lin_numerical(x))

    x_cat1 = f.gumbel_softmax(self.lin_cat_job(x), tau=0.2)
    x_cat2 = f.gumbel_softmax(self.lin_cat_sex(x), tau=0.2)
    x_cat3 = f.gumbel_softmax(self.lin_cat_incomeclass(x), tau=0.2)

    x_final = torch.cat((x_numerical, x_cat1, x_cat2, x_cat3),1)
    return x_final

I have managed to change the init part, using discrete_columns input which is an ordereddict that has the name and number of one-hot-encoding of each categorical field of my data as key and values, and continuous_columns which is only a list with the names of the continuous columns. But I have no idea how to edit the forward part:
class Generator(nn.Module):
  def __init__(self, input_dim, continuous_columns, discrete_columns):
    super(Generator, self).__init__()
    self._input_dim = input_dim
    self._discrete_columns = discrete_columns
    self._num_continuous_columns = len(continuous_columns)

    self.lin1 = nn.Linear(self._input_dim, self._input_dim)
    self.lin_numerical = nn.Linear(self._input_dim, self._num_continuous_columns)

    for key, value in self._discrete_columns.items():
      setattr(self, &quot;lin_cat_{}&quot;.format(key), nn.Linear(self._input_dim, value))
    
  def forward(self, x):
    x = torch.relu(self.lin1(x))
    x_numerical = f.leaky_relu(self.lin_numerical(x))
    ####
    This is the problematic part
    #####
    return x


You don't need to use setattr and honestly should not since you'd need getattr, it brings more trouble than it solves if there's any other ways to do the job.
Now this is what I'd do for this task
        self.lin_cat = nn.ModuleDict()
        for key, value in self._discrete_columns.items():
            self.lin_cat[key] = nn.Linear(self._input_dim, value)
        #     setattr(self, &quot;lin_cat_{}&quot;.format(key), nn.Linear(self._input_dim, value))

    def forward(self, x):
        x = torch.relu(self.lin1(x))
        x_numerical = f.leaky_relu(self.lin_numerical(x))

        x_cat = []
        for key in self.lin_cat:
            x_cat.append(f.gumbel_softmax(self.lin_cat[key](x), tau=0.2))
        x_final = torch.cat((x_numerical, *x_cat), 1)
        return x

