(self, dims: Sequence[int])
| 174 | |
| 175 | class MLP(nn.Sequential): |
| 176 | def __init__(self, dims: Sequence[int]): |
| 177 | nn.Sequential.__init__(self, |
| 178 | *itertools.chain(*[ |
| 179 | (nn.Linear(dim_in, dim_out), nn.ReLU(inplace=True)) |
| 180 | for dim_in, dim_out in zip(dims[:-2], dims[1:-1]) |
| 181 | ]), |
| 182 | nn.Linear(dims[-2], dims[-1]), |
| 183 | ) |
| 184 | |
| 185 | |
| 186 | class ConvStack(nn.Module): |