python - 从张量流中的张量返回k个最小元素

Question

我在张量流中有一个形状张量(16, 512, 4096)，我想k从张量中计算最小元素。

请注意，我可以使用以下代码片段在 pytorch 中获取它-

#inputs.shape (16L, 512L, 4096L)
dists, inputs_idx = torch.topk(inputs, 64, 2, largest=False, sorted=False)
#dists.shape (16L, 512L, 64L), inputs_idx.shape (16L, 512L, 64L)

请问有什么办法吗？

标签： pythontensorflow