import torch


def normalize_embeddings(embeddings, dim=0, order=2) -> torch.nn.Embedding:
    """Normalize embeddings in the specified dimension

    Parameters
    ----------
    embeddings : torch.nn.Embedding
    dim : int, optional
    order : float, optional

    Returns
    -------
    torch.nn.Parameter(torch.nn.Embedding) : the normalized embedding object

    """
    # norm = get_norm(embeddings.weight, dim=dim, order=order)
    norm = torch.norm(embeddings.weight, p=order, dim=dim, keepdim=True)
    normalized = torch.div(embeddings.weight, norm)
    embeddings.weight = torch.nn.Parameter(normalized)

    return embeddings