import torch


def apply_weight_norm(w, input_dims=(1, 2, 3), eps=1e-8):
    """
    Applies the "demodulation" operation from StyleGAN2 as a form of normalization.
    :param w: Weights
    :return: Normed weights
    """
    divisor = torch.sqrt((w ** 2).sum(dim=input_dims, keepdim=True) + eps)
    return w / divisor