(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs)
| 359 | |
| 360 | class ContextParallelCausalConv3d(nn.Module): |
| 361 | def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], stride=1, **kwargs): |
| 362 | super().__init__() |
| 363 | kernel_size = cast_tuple(kernel_size, 3) |
| 364 | |
| 365 | time_kernel_size, height_kernel_size, width_kernel_size = kernel_size |
| 366 | |
| 367 | assert is_odd(height_kernel_size) and is_odd(width_kernel_size) |
| 368 | |
| 369 | time_pad = time_kernel_size - 1 |
| 370 | height_pad = height_kernel_size // 2 |
| 371 | width_pad = width_kernel_size // 2 |
| 372 | |
| 373 | self.height_pad = height_pad |
| 374 | self.width_pad = width_pad |
| 375 | self.time_pad = time_pad |
| 376 | self.time_kernel_size = time_kernel_size |
| 377 | self.temporal_dim = 2 |
| 378 | |
| 379 | stride = (stride, stride, stride) |
| 380 | dilation = (1, 1, 1) |
| 381 | self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
| 382 | self.cache_padding = None |
| 383 | |
| 384 | def forward(self, input_, clear_cache=True): |
| 385 | # if input_.shape[2] == 1: # handle image |
nothing calls this directly
no test coverage detected