Hi,
I have a tensor a and want to select some data from it based on another tensor which is -1, 0, or 1. for example:
a = torch.randn([7, 4, 30])
b = tensor([[1, 0, 0, 0],
            [0, 1, 1, 1],
            [1, 0, 0, 0],
            [0, 1, 1, 1],
            [-1, 0, 0, 0],
            [-1,-1, 1, 1],
            [-1,-1, 0,-1]], dtype=torch.int32)

I want to select the data from each column of a where b==0 and b!=-1 or the same thing for b==1. I want to do it for each column and keep the data in a same column and then for example mask the rest of the column with zero or sth else. Finaly the shape of result would be [xx, 4, 30] where xx is sth less than 7 based on index values in that column. I tried to create mask for for example b==0 and use 'masked_select` but it gives me a 1-D tensor but I want the shape to be [xx, 4, 30].
Does anyone know a solution in Pytorch? Number of columns and rows are big in mycode and I don’t want to use for loop to do this.
Update:
a = torch.randn([6, 4, 30])
b = tensor([[1, 0, 0, 0],
            [0, 1, 1, 1],
            [1, 0, 0, 0],
            [0, 1, 1, 1],
            [-1, 0, 0, 0],
            [-1,-1, 1, 1]], dtype=torch.int32)