Method __init__

sat/sgm/modules/autoencoding/magvit2_pytorch.py:454–466

(self, dim, *, dim_cond: Optional[int] = None, mult=4, images=False)

class FeedForward(Module):
    @beartype
    def __init__(self, dim, *, dim_cond: Optional[int] = None, mult=4, images=False):
        super().__init__()
        conv_klass = nn.Conv2d if images else nn.Conv3d

        rmsnorm_klass = RMSNorm if not exists(dim_cond) else partial(AdaptiveRMSNorm, dim_cond=dim_cond)

        maybe_adaptive_norm_klass = partial(rmsnorm_klass, channel_first=True, images=images)

        dim_inner = int(dim * mult * 2 / 3)

        self.norm = maybe_adaptive_norm_klass(dim)

        self.net = Sequential(conv_klass(dim, dim_inner * 2, 1), GEGLU(), conv_klass(dim_inner, dim, 1))

    @beartype
    def forward(self, x: Tensor, *, cond: Optional[Tensor] = None):