Slicing 3d tensor with 2d tensors

1

Hi,
I have 3d tensor like this:

tensor([[[0.6066, 0.0390, 0.0141, 0.3977, 0.2390, 0.1982, 0.8321],
[0.5124, 0.5975, 0.6119, 0.0933, 0.5048, 0.6961, 0.4293],
[0.0222, 0.7371, 0.0867, 0.0660, 0.5737, 0.4022, 0.9751],
[0.6079, 0.9880, 0.5903, 0.8366, 0.9752, 0.1610, 0.8460],
[0.3080, 0.8679, 0.4383, 0.2642, 0.7401, 0.6404, 0.0325]],
[[0.9429, 0.7948, 0.9669, 0.7277, 0.9147, 0.9281, 0.4738],
[0.0109, 0.7387, 0.5509, 0.5736, 0.3486, 0.9945, 0.6516],
[0.0915, 0.0928, 0.2610, 0.0317, 0.3806, 0.2865, 0.5693],
[0.1901, 0.1824, 0.1148, 0.0235, 0.1220, 0.3003, 0.7579],
[0.8647, 0.9720, 0.6833, 0.6454, 0.0365, 0.6508, 0.5220]]])

and 2d tensor like this:

tensor([[1, 3], [2, 3]])

I want to get tensor in range (1, 2, 3), (2, 3) based on index tensor to get output:

tensor([[[0.5124, 0.5975, 0.6119, 0.0933, 0.5048, 0.6961, 0.4293],
[0.0222, 0.7371, 0.0867, 0.0660, 0.5737, 0.4022, 0.9751],
[0.6079, 0.9880, 0.5903, 0.8366, 0.9752, 0.1610, 0.8460]],
[[0.0915, 0.0928, 0.2610, 0.0317, 0.3806, 0.2865, 0.5693],
[0.1901, 0.1824, 0.1148, 0.0235, 0.1220, 0.3003, 0.7579]]])

then average tensor by batch.

Any idea or help on how to solve this is well appreciated!
Thank you in advance.

2

I don’t quite understand the second tensor shape and how it corresponds to the index tensor.
However, since the index tensor is using a different number of indices, you would need to either transform the indices to a flattened version (and the same with the original input) or index them separately.

3

Do you mean that you want to use the second tensor for indexing the first tensor? Did you try torch.gather?