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大的数,然后在原矩阵中找其位置,获取位置信息。
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大的数,然后在原矩阵中找其位置,获取位置信息。