Slice tensor using boolean tensor

M = torch.randn(4,2)
idx_mask = torch.Tensor([1,1,0,0]).long().view(-1,1)

I want to use idx_mask in order to get rows 0&1 from m.
This operation is done easily in MATLAB by M(idx_mask, :).
What is the best way to do it in PyTorch ?

i know i can expand idx_mask to have the same dimensions as M and do the slicing, but i’m curious to know if there is a better way.

1 Like

you can use the nonzero function to get indices from a mask:

M = torch.randn(4,2)
idx_mask = torch.Tensor([1,1,0,0]).long()
M[idx_mask.nonzero(), :]
2 Likes

Thanks for this elegant solution!

Is this .zonzero() function differentiable?

it’s not differentiable because it is a discrete operation