r""" Image to Patch Unembedding Args: embed_dim (int): Number of linear projection output channels.
| 339 | |
| 340 | |
| 341 | class PatchUnEmbedIR(nn.Module): |
| 342 | r""" Image to Patch Unembedding |
| 343 | |
| 344 | Args: |
| 345 | embed_dim (int): Number of linear projection output channels. |
| 346 | """ |
| 347 | |
| 348 | def __init__(self, embed_dim=96): |
| 349 | super().__init__() |
| 350 | self.embed_dim = embed_dim |
| 351 | |
| 352 | def forward(self, x, x_size): |
| 353 | B, H, W, C = x.shape |
| 354 | x = x.permute(0, 3, 1, 2) |
| 355 | x = x.view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C |
| 356 | return x |