(
self,
f_channels: int,
zq_channels: int,
)
| 4188 | """ |
| 4189 | |
| 4190 | def __init__( |
| 4191 | self, |
| 4192 | f_channels: int, |
| 4193 | zq_channels: int, |
| 4194 | ): |
| 4195 | super().__init__() |
| 4196 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 4197 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 4198 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 4199 | |
| 4200 | def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: |
| 4201 | f_size = f.shape[-2:] |