import torch


def pi_to_y(Pi):
    """
    Turns a class membership matrix into a vector of class labels.
    Args:
        Pi: A one-hot encoded class membership matrix (N, K)

    Returns:
        A vector y of class labels (N, )
    """
    return torch.argmax(
        Pi,  # (N, K)
        dim=1
    )