Coming here having a similar question. If I understood the OP correctly, given a tensor A of
(1, 3, 5) and a corresponding tensor B of shape (5,) where for example
A=
[[[3, 0, 0, 1, 8, 9, 5, 4, 1, 6],
[0, 0, 0, 8, 0, 7, 1, 0, 0, 8],
[1, 0, 2, 5, 3, 6, 3, 2, 4, 7]]]
and B=[0, 0, 1, 1, 0, 2, 2, 3, 3, 3]
the goal is to get 4 tensors where the first contains all elements of A whose corresponding B is 0.
Thus
[[[3, 0, 8],
[0, 0, 0],
[1, 0, 3]]]
the second one would be
[[[0, 1],
[0, 8],
[2, 5]]]
and so forth.
Is there an easy was of doing it. And to answer my actual question, is there also a way to generalize it to arbitrary index values and then split in a similar fashion such that each resulting tensor contains values with the same corresponding index?