以二维numpy矩阵为例
import torch
import numpy as np
K=3 #取每行最小3个值的索引
data=np.random.rand(4,7)
print(data)
data=torch.from_numpy(data)
a, idx = torch.sort(data, descending=False)
lists=idx[:,:K]
print(lists)
import torch
import numpy as np
K=3 #取每行最小3个值的索引
data=np.random.rand(4,7)
print(data)
data=torch.from_numpy(data)
a, idx = torch.sort(data, descending=False)
lists=idx[:,:K]
print(lists)