| 542 | |
| 543 | |
| 544 | class Downsample3D(nn.Module): |
| 545 | def __init__(self, |
| 546 | in_channels, |
| 547 | with_conv, |
| 548 | stride |
| 549 | ): |
| 550 | super().__init__() |
| 551 | |
| 552 | self.with_conv = with_conv |
| 553 | if with_conv: |
| 554 | self.conv = CausalConv(in_channels, in_channels, kernel_size=3, stride=stride) |
| 555 | |
| 556 | def forward(self, x, is_init=True): |
| 557 | if self.with_conv: |
| 558 | x = self.conv(x, is_init) |
| 559 | else: |
| 560 | x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) |
| 561 | return x |
| 562 | |
| 563 | class VideoEncoder(nn.Module): |
| 564 | def __init__(self, |