I have a tensor of order 5 and want to reduce it to order 4 by selectively picking along the third domain based on the position in the second domain.

Let’s say my original tensor’s shape is `(1, 16, 16, 8, 8)`

and I want to get `(1, 16, 8, 8)`

. So far, I am doing this as follows iteratively:

```
import torch
original_tensor = torch.randint(10, (1, 4, 4, 2, 2))
output_tensor = torch.empty((1, 4, 2, 2))
for i in range(output_tensor.shape[1]):
output_tensor[:, i, :, :] = original_tensor[:, i, i, :, :].squeeze(1)
```

Putting the outputs in here would make the post very long, but you can execute the snippet on its own and print out both `original_tensor`

and `output_tensor`

. However for a second order into first order tensor consider this example:

```
ORIGINAL_TENSOR
tensor([[5, 8, 2, 5],
[1, 4, 5, 8],
[1, 5, 6, 0],
[3, 2, 7, 2]])
OUTPUT_TENSOR
tensor([5., 4., 6., 2.])
```

Iteratively, as you can imagine, this is way to slow given that as part of a layer I need to execute this a lot. I do though beforehand know (that is, it is constant) the indices which i want to keep. So I could precalculate some sort of mask if that simplifies it.

Please consider the added complexity due to the order and size of the real example I gave in the first snippet. A solution that just helps with the particular second example would not help.