| 1684 | """ |
| 1685 | |
| 1686 | def __init__( |
| 1687 | self, |
| 1688 | f_channels: int, |
| 1689 | zq_channels: int, |
| 1690 | ): |
| 1691 | super().__init__() |
| 1692 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 1693 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1694 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1695 | |
| 1696 | def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor: |
| 1697 | f_size = f.shape[-2:] |