Выберите определенные c строки 2D тензора PyTorch

Question

Предположим, у меня есть 2D-тензор, который выглядит примерно так:

[[44, 50, 1, 32],
.
.
.
[7, 13, 90, 83]]

и список индексов строк, которые я хочу выбрать, выглядит примерно так [0, 34, 100, ..., 745]. Как я могу go пройти и создать новый тензор, содержащий только те строки, индексы которых содержатся в массиве?

Answer 1

Вы можете выбрать как с numpy

import torch
x = torch.Tensor([[1, 2, 3, 4],
                  [5, 6, 7, 8],
                  [9, 8, 7, 6],
                  [5, 4, 2, 1]])

indices = [0, 3]

print(x[indices])
# tensor([[1., 2., 3., 4.],
#         [5., 4., 2., 1.]])