pytorch - 在 PyTorch 中用张量索引多维张量

Question

我有以下代码：

a = torch.randint(0,10,[3,3,3,3])
b = torch.LongTensor([1,1,1,1])

我有一个多维索引b，想用它来选择a. 如果 b 不是张量，我可以这样做：

a[1,1,1,1]

返回正确的单元格，但是：

a[b]

不起作用，因为它只选择a[1]了四次。

我怎样才能做到这一点？谢谢

标签： pytorchtensor

Answer 1

一个更优雅（更简单）的解决方案可能是简单地转换b为一个元组：

a[tuple(b)]
Out[10]: tensor(5.)

我很想知道它是如何与“常规”numpy 一起工作的，并在这里找到了一篇相关文章很好地解释了这一点。