| 225 | |
| 226 | |
| 227 | class ImageEmbedding(nn.Module): |
| 228 | def __init__(self, in_channels=7, out_channels=320) -> None: |
| 229 | super().__init__() |
| 230 | self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) |
| 231 | |
| 232 | def forward(self, x) -> torch.Tensor: |
| 233 | return self.f(x) |
| 234 | |
| 235 | |
| 236 | class ImageUnembedding(nn.Module): |