Let’s consider
a = torch.arange(8).view(2,2,2)

which should look like
[[[0, 1],
[2, 3]],

[[4, 5],
[6, 7]]]

and also
index = torch.tensor([1, 0])

“index” shows column indexes which should be extracted from “a”.
so the output should be look like this:
[[1, 4],
 [3, 6]]

although it seems not difficult, using torch.gather and other indexing methods couldn’t help me. Any suggestions?