Function l2norm

model/layers.py:7–13

L2-normalize columns of X

(X, dim=-1, eps=1e-12)

from timm.models.layers import trunc_normal_

def l2norm(X, dim=-1, eps=1e-12):
    """
    L2-normalize columns of X
    """
    norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps
    X = torch.div(X, norm)
    return X


class Mlp(nn.Module):