def get_id_from_tensor(tensor, id_name):
    """Gets appropriate id column from tensor with ids"""
    if id_name == "user":
        return tensor[:, 0]

    elif id_name == "item":
        return tensor[:, 1]

    else:
        raise ValueError(f"id_name should be one of 'user', 'item', got {id_name}")