How can I efficiently map values from a 3D tensor to a 1D tensor using indices in PyTorch?

I'm trying to replace each element in a 3D tensor (shape: i, j, j) with a corresponding value from a 1D tensor (shape: n) by using the values in the 3D tensor as indices into the 1D tensor.

I have two tensors:

• A 3D tensor (shape: i, j, j) with integer values ranging from 0 to n
• A 1D tensor (shape: n) containing corresponding values

I need to create a result tensor with the same shape as the 3D tensor, where each element is replaced by the value at that index from the 1D tensor (e.g., if a certain element in the 3D tensor is 42, it should be replaced with the value at index 42 in the 1D tensor).

Here’s an example demonstrating the desired behavior:

PYTHONCopied
import torch

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

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

# How it should work
print(big_tensor[0, 0, 0])    # Suppose this outputs 42
print(skip_tensor[42])       # This might output 0.7853
print(result[0, 0, 0])        # This should also be 0.7853

I want to achieve this without using explicit loops for performance reasons, as both tensors may be quite large. Is there an efficient built-in PyTorch operation or method for this task?