import torch


def topk_correct(pred, target, topk):
    """Computes the accuracy of the correct top-k predictions for the specified values of k"""

    assert pred.size(0) == target.size(0), "Batch dim of predictions and targets must match"
    with torch.no_grad():
        max_k = max(topk)

        top_max_k_vals, top_max_k_inds = torch.topk(pred, max_k, dim=1, largest=True, sorted=True)

        # (batch_size, max_k) -> (max_k, batch_size)
        top_max_k_inds = top_max_k_inds.t()
        # (batch_size, ) -> (max_k, batch_size)
        rep_max_k_lbls = target.view(1, -1).expand_as(top_max_k_inds)

        top_max_k_correct = top_max_k_inds.eq(rep_max_k_lbls)

        # compute the number of top-k correct predictions
        topk_correct_list = []
        for k in topk:
            correct_k = top_max_k_correct[:k, :].view(-1).float().sum()
            topk_correct_list.append(correct_k)

        return topk_correct_list