sajadn

import torch.nn as nn
from torch import _weight_norm
from torch.nn.utils import weight_norm


def data_dependent_init(model, data_batch):
    def init_hook_(module, input, output):
        std, mean = torch.std_mean(output, dim=[0, 2, 3])
        g = getattr(module, 'weight_g')
        g.data = g.data/std.reshape((len(std), 1, 1, 1))