Function ResidualUnit

sat/sgm/modules/autoencoding/magvit2_pytorch.py:843–852

(dim, kernel_size: Union[int, Tuple[int, int, int]], pad_mode: str = "constant")

@beartype
def ResidualUnit(dim, kernel_size: Union[int, Tuple[int, int, int]], pad_mode: str = "constant"):
    net = Sequential(
        CausalConv3d(dim, dim, kernel_size, pad_mode=pad_mode),
        nn.ELU(),
        nn.Conv3d(dim, dim, 1),
        nn.ELU(),
        SqueezeExcite(dim),
    )

    return Residual(net)


@beartype