import torch


def unit_vec(vec, dim=-1):
    """
    Returns a normalized unit vector along the provided dimension

    Args:
        vec (tensor): (..., N, ...) arbitrary vector
        dim (int): Dimension to normalize

    Returns:
        tensor: (..., N, ...) vector normalized along @dim
    """
    # type: (Tensor, int) -> Tensor
    return vec / torch.norm(vec, dim=dim, keepdim=True)