Function Normalize

sat/vae_modules/cp_enc_dec.py:444–448

(in_channels, gather=False, **kwargs)

def Normalize(in_channels, gather=False, **kwargs):  # same for 3D and 2D
    if gather:
        return ContextParallelGroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
    else:
        return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)


class SpatialNorm3D(nn.Module):