| 71 | |
| 72 | |
| 73 | class Encoder(nn.Module): |
| 74 | def __init__(self, config): |
| 75 | super().__init__() |
| 76 | self.config = config |
| 77 | self.blocks = self.init_blocks(config) |
| 78 | |
| 79 | def init_blocks(self, config): |
| 80 | blocks = [] |
| 81 | for block in config['blocks']: |
| 82 | blocks.append(Block(**block)) |
| 83 | return nn.Sequential(*blocks) |
| 84 | |
| 85 | def forward(self, x): |
| 86 | return self.blocks(x) |
| 87 | |
| 88 | |
| 89 | class Decoder(nn.Module): |