(self, dim_in, dim_out, dim_hidden, num_layers)
| 68 | |
| 69 | class MLP(nn.Module): |
| 70 | def __init__(self, dim_in, dim_out, dim_hidden, num_layers): |
| 71 | super().__init__() |
| 72 | self.dim_in = dim_in |
| 73 | self.dim_out = dim_out |
| 74 | self.dim_hidden = dim_hidden |
| 75 | self.num_layers = num_layers |
| 76 | |
| 77 | net = [] |
| 78 | for l in range(num_layers): |
| 79 | net.append(nn.Linear(self.dim_in if l == 0 else self.dim_hidden, self.dim_out if l == num_layers - 1 else self.dim_hidden, bias=False)) |
| 80 | |
| 81 | self.net = nn.ModuleList(net) |
| 82 | |
| 83 | def forward(self, x): |
| 84 | for l in range(self.num_layers): |