import torch


def categorical_to_topk(predictions, k=1):
    """ Convert a categorical prediction to a a top-k prediction by taking the k indices with largest probability
    
    This function is not differentiable. 
    
    Args:
        predictions (tensor): a batch of categorical predictions.
        k (int): the number of predicted labels.
    
    Returns:
        tensor: the topk prediction. 
    """
    topk = torch.topk(predictions, k, dim=1)[1]
    if k == 1:
        topk = topk.flatten()
    return topk