| 290 | |
| 291 | |
| 292 | class UVit2DConvEmbed(nn.Module): |
| 293 | def __init__(self, in_channels, block_out_channels, vocab_size, elementwise_affine, eps, bias): |
| 294 | super().__init__() |
| 295 | self.embeddings = nn.Embedding(vocab_size, in_channels) |
| 296 | self.layer_norm = RMSNorm(in_channels, eps, elementwise_affine) |
| 297 | self.conv = nn.Conv2d(in_channels, block_out_channels, kernel_size=1, bias=bias) |
| 298 | |
| 299 | def forward(self, input_ids): |
| 300 | embeddings = self.embeddings(input_ids) |
| 301 | embeddings = self.layer_norm(embeddings) |
| 302 | embeddings = embeddings.permute(0, 3, 1, 2) |
| 303 | embeddings = self.conv(embeddings) |
| 304 | return embeddings |
| 305 | |
| 306 | |
| 307 | class UVitBlock(nn.Module): |