- layers = [nn.Linear(2, nh), activation()]
- nb_hidden = 4
- for k in range(nb_hidden):
- layers += [nn.Linear(nh, nh), activation()]
- layers += [nn.Linear(nh, 2)]
+ layers = [nn.Linear(2, hidden_dim), activation()]
+ for k in range(nb_hidden - 1):
+ layers += [nn.Linear(hidden_dim, hidden_dim), activation()]
+ layers += [nn.Linear(hidden_dim, 2)]