Function Normalize3D

sat/vae_modules/cp_enc_dec.py:513–528

(
    in_channels,
    zq_ch,
    add_conv,
    gather=False,
)

def Normalize3D(
    in_channels,
    zq_ch,
    add_conv,
    gather=False,
):
    return SpatialNorm3D(
        in_channels,
        zq_ch,
        gather=gather,
        freeze_norm_layer=False,
        add_conv=add_conv,
        num_groups=32,
        eps=1e-6,
        affine=True,
    )


class Upsample3D(nn.Module):