Forces tensors to have norms within a range

max_norm.py

class MaxNorm(object):
    def __init__(self, max_value=1, frequency=5):
        self.frequency = frequency
        self.max_value = max_value
        self.tiny = _finfo(torch.FloatTensor([])).tiny

    def __call__(self, module):
        if hasattr(module, "weight"):
            w = module.weight.data
            norms = w.norm(p=2, dim=w.dim() - 1, keepdim=True)
            desired = norms.clamp(0, self.max_value)
            w *= desired / (self.tiny + norms)
            return w

        if hasattr(module, "bias"):
            b = module.bias.data
            norms = b.norm(p=2, dim=b.dim() - 1, keepdim=True)
            desired = norms.clamp(0, self.max_value)
            b *= desired / (self.tiny + norms)
            return b