| 412 | |
| 413 | |
| 414 | class CausalConvChannelLast(CausalConv): |
| 415 | def __init__(self, |
| 416 | chan_in, |
| 417 | chan_out, |
| 418 | kernel_size, |
| 419 | **kwargs |
| 420 | ): |
| 421 | super().__init__( |
| 422 | chan_in, chan_out, kernel_size, **kwargs) |
| 423 | |
| 424 | self.time_causal_padding = (0, 0) + self.time_causal_padding |
| 425 | self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding |
| 426 | |
| 427 | def forward(self, x, is_init=True, residual=None): |
| 428 | if self.is_first_run: |
| 429 | self.is_first_run = False |
| 430 | # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) |
| 431 | |
| 432 | x = nn.functional.pad(x, |
| 433 | self.time_causal_padding if is_init else self.time_uncausal_padding) |
| 434 | |
| 435 | x = base_conv3d_channel_last(x, self.conv, residual=residual) |
| 436 | return x |
| 437 | |
| 438 | class CausalConvAfterNorm(CausalConv): |
| 439 | def __init__(self, |