Tensor Manipulatation

Hi, I’m a beginner of PyTorch and I wonder how to manipulate tensors. Suppose that there are a LongTensor a = torch.LongTensor([1, 2, 3, 4, 5, -1, -1, -1, -1]) of shape (9,) and a FloatTensor b = torch.FloatTensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8], [9, 9]]) of shape (9, 2). Then I want to drop those values equal to -1, so I set a_drop = a[a != -1] and I get [1, 2, 3, 4, 5] successfully. But when I want to get corresponding b_drop, say, [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]] and set b_drop = b[a != -1] but I failed.

So I wonder how to use a != 1 as a mask to get corresponding b_drop. Maybe there are better methods, please told me anyway. Thanks!

You could use:

b_drop = b[(a != -1).nonzero().squeeze()]