Как отобразить значения от трехмерного тензора до 1D тензора в Pytorch? ⇐ Python

[Anonymous]

Я застрял с проблемой Pytorch, и мне нужна помощь:
У меня есть два тензора:

[*]3D-тензор (форма: i, j, j) с целыми значениями от 0 до n
[*]1D-тензор (форма: n)
< /ul>
Мне нужно создать новый тензор это та же форма, что и первая (i, j, j), но каждое значение заменяется соответствующим значением из второго тензора. Это похоже на использование значений первого тензора в качестве индексов для второго.
Вот краткий пример:

Код: Выделить всё

import torch

# My tensors
big_tensor = torch.randint(0, 256, (10, 25, 25))
small_tensor = torch.rand(256)

# What I'm trying to do
result = magic_function(big_tensor, small_tensor)

# How it should work
print(big_tensor[0, 0, 0])    # Let's say this outputs 42
print(small_tensor[42])       # This might output 0.7853
print(result[0, 0, 0])        # This should also be 0.7853

Я ищу эффективный способ сделать это, желательно без использования циклов, поскольку оба тензора могут быть довольно большими. Есть ли эффективная операция или метод Pytorch, который я могу использовать для этого?
Любая помощь будет замечательной — заранее спасибо!

Подробнее здесь: https://stackoverflow.com/questions/793 ... in-pytorch