(dim_list, latent_dim)
| 2 | |
| 3 | |
| 4 | def build_MLP(dim_list, latent_dim): |
| 5 | model_list = [] |
| 6 | prev = dim_list[0] |
| 7 | for cur in dim_list[1:]: |
| 8 | model_list.append(nn.Linear(prev, cur)) |
| 9 | model_list.append(nn.GELU()) |
| 10 | prev = cur |
| 11 | model_list.append(nn.Linear(prev, latent_dim)) |
| 12 | model = nn.Sequential(*model_list) |
| 13 | return model |