| 628 | class Conv3DMod(Module): |
| 629 | @beartype |
| 630 | def __init__( |
| 631 | self, dim, *, spatial_kernel, time_kernel, causal=True, dim_out=None, demod=True, eps=1e-8, pad_mode="zeros" |
| 632 | ): |
| 633 | super().__init__() |
| 634 | dim_out = default(dim_out, dim) |
| 635 | |
| 636 | self.eps = eps |
| 637 | |
| 638 | assert is_odd(spatial_kernel) and is_odd(time_kernel) |
| 639 | |
| 640 | self.spatial_kernel = spatial_kernel |
| 641 | self.time_kernel = time_kernel |
| 642 | |
| 643 | time_padding = (time_kernel - 1, 0) if causal else ((time_kernel // 2,) * 2) |
| 644 | |
| 645 | self.pad_mode = pad_mode |
| 646 | self.padding = (*((spatial_kernel // 2,) * 4), *time_padding) |
| 647 | self.weights = nn.Parameter(torch.randn((dim_out, dim, time_kernel, spatial_kernel, spatial_kernel))) |
| 648 | |
| 649 | self.demod = demod |
| 650 | |
| 651 | nn.init.kaiming_normal_(self.weights, a=0, mode="fan_in", nonlinearity="selu") |
| 652 | |
| 653 | @beartype |
| 654 | def forward(self, fmap, cond: Tensor): |