【发布时间】:2020-12-07 08:05:09
【问题描述】:
我正在使用以下代码使用 pytorch 查找 topk 匹配项:
def find_top(self, x, y, n_neighbors, unit_vectors=False, cuda=False):
if not unit_vectors:
x = __to_unit_torch__(x, cuda=cuda)
y = __to_unit_torch__(y, cuda=cuda)
with torch.no_grad():
d = 1. - torch.matmul(x, y.transpose(0, 1))
values, indices = torch.topk(d, n_neighbors, dim=1, largest=False, sorted=True)
return indices.cpu().numpy()
不幸的是,它抛出以下错误:
values, indices = torch.topk(d, n_neighbors, dim=1, largest=False, sorted=True)
RuntimeError: invalid argument 5: k not in range for dimension at /pytorch/aten/src/THC/generic/THCTensorTopK.cu:23
d 的大小是 (1793,1) 。我错过了什么?
【问题讨论】:
标签: python python-3.x pytorch top-n