Индексирование в двумерном тензоре PyTorch с использованием другого тензора

Question

Предположим, что тензор A определен как:

 1  2  3  4
 5  6  7  8
 9 10 11 12
13 14 15 16

Я пытаюсь извлечь плоский массив из этой матрицы, используя другой тензор в качестве индексов. Например, если второй тензор определен как:

0
1
2
3

Я хочу, чтобы результатом индексации был 1-D тензор с содержанием:

1
6
11
16

Это не похоже, ведут себя как NumPy; Я попытался A[:, B], но он просто выдает ошибку из-за невозможности выделить безумное количество памяти, и я понятия не имею, почему!

Answer 1

Вы можете преобразовать свой Тензор в массив NumPy. Если вы используете Cuda, не забудьте передать его процессору. Если нет, то нет необходимости передавать его процессору. Пример кода ниже:

val.data.cpu().numpy()[:,B]

Дайте мне знать, если это решит вашу проблему

Answer 2

PyTorch реализует torch.take, который эквивалентен numpy .take