import torch


def l2_norm(input, axis=1):
    """l2 normalization.

    Args:
        input (torch.Tensor): The input tensor.
        axis (int, optional): Specifies which axis of input to calculate the
            norm across. Defaults to 1.

    Returns:
        Tensor: Tensor after L2 normalization per-instance.
    """
    norm = torch.norm(input, 2, axis, True)
    output = torch.div(input, norm)
    return output