import torch


def l2n(x: torch.Tensor, eps: float = 1e-6, dim: int = 1):
    """
    L2 normalize the input tensor along the specified dimension

    Args:
        x (torch.Tensor): the tensor to normalize
        eps (float): epsilon to use to normalize to avoid the inf output
        dim (int): along which dimension to L2 normalize

    Returns:
        x (torch.Tensor): L2 normalized tensor
    """
    x = x / (torch.norm(x, p=2, dim=dim, keepdim=True) + eps).expand_as(x)
    return x