r""" Image to Patch Embedding Args: embed_dim (int): Number of linear projection output channels. norm_layer (nn.Module, optional): Normalization layer.
| 319 | return mask_h, mask_w |
| 320 | |
| 321 | class PatchEmbedIR(nn.Module): |
| 322 | r""" Image to Patch Embedding |
| 323 | |
| 324 | Args: |
| 325 | embed_dim (int): Number of linear projection output channels. |
| 326 | norm_layer (nn.Module, optional): Normalization layer. |
| 327 | """ |
| 328 | |
| 329 | def __init__(self, embed_dim=96, norm_layer=None): |
| 330 | super().__init__() |
| 331 | self.norm = nn.LayerNorm(embed_dim) if norm_layer is not None else None |
| 332 | |
| 333 | def forward(self, x): |
| 334 | x = x.permute(0, 2, 3, 1) # (b c h w) -> (b h w c) |
| 335 | if self.norm is not None: |
| 336 | # print("Using norm layer") |
| 337 | x = self.norm(x) |
| 338 | return x |
| 339 | |
| 340 | |
| 341 | class PatchUnEmbedIR(nn.Module): |