Function Normalize

diffbir/model/vae.py:18–21

(in_channels, num_groups=32)

def Normalize(in_channels, num_groups=32):
    return torch.nn.GroupNorm(
        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
    )


class Upsample(nn.Module):