Softmax中top K个元素

afterarr, _ = torch.sort(template_pred, descending=True)
kdata = afterarr[:, k]
kindex = (template_pred[0] == kdata).nonzero()
kindex = kindex.cpu().tolist()[0]
template_id = kindex[0]

先对原来数组排序，然后取出第K大的数，然后在原矩阵中找其位置，获取位置信息。