(self, config)
| 48 | return init_embedding(config) |
| 49 | |
| 50 | def init_encoder(self, config): |
| 51 | self.encoder = Encoder(config) |
| 52 | # Allow access to first encoder SSM kernel_dim |
| 53 | self.kernel_dim = self.encoder.blocks[0].ssm.kernel_dim |
| 54 | return self.encoder |
| 55 | |
| 56 | def init_decoder(self, config): |
| 57 | self.decoder = Decoder(config) |