I'm stuck with a Pytorch problem and could use some help: I've got two tensors:  A 3D tensor (shape: i, j, j) with integer values from 0 to n A 1D tensor (shape: n)  I need to create a new tensor that's the same shape as the first one (i, j, j), but where each value is replaced by the corresponding value from the second tensor. Like using the values in the first tensor as indices for the second one. Here's a quick example: 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  I'm looking for a performant way to do this, preferably not using loops, as both tensors can be quite big. Is there an efficient Pytorch operation or method I can use for this? Any help would be awesome - thanks in advance!
This should work: small_tensor[big_tensor]  Take note that the type of the big_tensor must be long/int.  Edit: In response to the comment of @simon, I wrote a colab notebook that shows how this solution works without the need to perform any other operation.