Retrieving values from a 5D tensor using indices calculated via nonzero()

1

Hi everyone,
I am a little bit confused on how I could use torch.gather() or torch.index_select() to retrieve some specific values from a 5D tensor using a 2D tensor of indices I obtained via torch.nonzero(). Say I have two 5d tensors A and B, B being extremely sparse (95% of zeros):

>> A.size()
torch.Size([20, 20, 64, 64, 64])
>>B.size()
torch.Size([20, 20, 64, 64, 64])

If I call nonzero() on B, I get a 2D tensor with for each row the 5 indices of a nonzero value in B:

>> ind = B.nonzero()
>> print(ind)
 tensor([[ 0,  1, 28, 32, 34],
        [ 0,  1, 32, 35, 39],
        [ 0,  2, 37, 26, 45],
        ...,
        [19, 17, 20, 39, 37],
        [19, 18, 27, 38, 31],
        [19, 19, 23, 38, 36]]
)

How can I get the corresponding values in A using ind ?

2

OK, I found a way to circumvent this problem using torch. masked_select . Thanks!