(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs)
| 294 | |
| 295 | class ContextParallelCausalConv3d(nn.Module): |
| 296 | def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs): |
| 297 | super().__init__() |
| 298 | kernel_size = cast_tuple(kernel_size, 3) |
| 299 | |
| 300 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 301 | |
| 302 | assert is_odd(height_kernel_size) and is_odd(width_kernel_size) |
| 303 | |
| 304 | time_pad = time_kernel_size - 1 |
| 305 | height_pad = height_kernel_size // 2 |
| 306 | width_pad = width_kernel_size // 2 |
| 307 | |
| 308 | self.height_pad = height_pad |
| 309 | self.width_pad = width_pad |
| 310 | self.time_pad = time_pad |
| 311 | self.time_kernel_size = time_kernel_size |
| 312 | self.temporal_dim = 2 |
| 313 | |
| 314 | stride = (stride, stride, stride) |
| 315 | dilation = (1, 1, 1) |
| 316 | self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
| 317 | |
| 318 | def forward(self, input_): |
| 319 | # temporal padding inside |
nothing calls this directly
no test coverage detected