Converts a 1D tensor of predicted labels into a 2D tensor of probabilistic labels

    Args:
        Y_h: an [n], or [n,1] tensor of predicted (int) labels in {1,...,k}
        k: the largest possible label in Y_h
    Returns:
        Y_s: a torch.FloatTensor of shape [n, k] where Y_s[i, j-1] is the probabilistic
            label for item i and label j

def pred_to_prob(Y_h, k):
    """Converts a 1D tensor of predicted labels into a 2D tensor of probabilistic labels

    Args:
        Y_h: an [n], or [n,1] tensor of predicted (int) labels in {1,...,k}
        k: the largest possible label in Y_h
    Returns:
        Y_s: a torch.FloatTensor of shape [n, k] where Y_s[i, j-1] is the probabilistic
            label for item i and label j
    """
    Y_h = Y_h.clone()
    if Y_h.dim() > 1:
        Y_h = Y_h.squeeze()
    assert Y_h.dim() == 1
    assert (Y_h >= 1).all()
    assert (Y_h <= k).all()
    n = Y_h.shape[0]
    Y_s = torch.zeros((n, k), dtype=Y_h.dtype, device=Y_h.device)
    for i, j in enumerate(Y_h):
        Y_s[i, j - 1] = 1.0
    return Y_s