Olá Vinícius,
Você pode fazer algo assim:
class classificador_torch(nn.Module):
def __init__(self, activation, neurons, n_hidden_layers, initializer):
super().__init__()
self.input_layer = nn.Linear(30, neurons)
initializer(self.input_layer.weight)
self.hidden_layers = [nn.Linear(neurons, neurons) for _ in range(n_hidden_layers)]
for layer in self.hidden_layers:
initializer(layer.weight)
self.activation = activation
self.output_layer = nn.Linear(neurons, 1)
def forward(self, X):
X = self.input_layer(X)
X = self.activation(X)
for layer in self.hidden_layers:
X = layer(X)
X = self.activation(X)
X = self.output_layer(X)
return X