import torch


def get_prediction_ids(predictions):
    """ Get the label ids from the raw prediction outputs """
    tensor = torch.tensor(predictions)
    softmax = torch.nn.functional.softmax(tensor, dim=-1)
    argmax = torch.argmax(softmax, dim=-1)
    return argmax.tolist()