根据索引沿维度选择张量切片
pytorch 186
原文标题 :Select tensor slice along a dimension based on index
我有以下形状的 PyTorch 张量:(100, 5, 100)。我需要通过从每一行中仅选择第二维中的一项来将其转换为形状为 (100, 100) 的张量,这意味着在这 5 个元素中我只需要一个,以及其对应的 100 个元素。
要执行此操作,我有第二个形状为 (100,) 的张量,其索引指定应在每行中选择这 5 个项目中的哪一个。
有没有一种简单的方法来执行此选择而不必过多地弄乱尺寸?