import torch


def topk(A, k, dim=1):
    """Returns the k largest elements of the given input tensor along a given dimension.

    Parameters
    ----------
    A
    k
    dim

    Returns
    -------
    values : tensor of shape (n_samples, k)
        Top k values.

    index : tensor of shape (n_samples, k)
        Top k indexes.

    """
    if len(A.shape) == 1:
        dim = 0
    tk = torch.topk(A.cuda(), k, dim=dim)
    # print(A.cuda())
    return tk[0].cpu(), tk[1].cpu()