import torch


def normalize_vec(Z):  # (N, D)
    """
    Normalizes each row of the tensor Z by dividing it by its Frobenius (l2) norm.
    Args:
        Z: the data tensor (N, D)

    Returns:
        A normalized version of Z (N, D)
    """
    return torch.nn.functional.normalize(Z, dim=1)