Using torch.nonzero to index into a larger tensor

1

Relatively new to pytorch so this might already be answered. Suppose I have a tensor y of shape (H, W) so that

idx = torch.nonzero(y)

has shape (N, 2) where N is the number of nonzero entries of y. I also have a tensor x of shape (4, H, W) and I want to create an output tensor z of shape (N, 4) where I’ve used idx to index into the last two dimensions of x. In my code I want to return idx and z as parallel tensors. There has to be a neat way to do this but I’m not there yet with pytorch.

Thanks,
Dave

2

Hi Dave,

You could have a try on the code snippet as follows, I’m not sure whether it is efficient, but it works.

z = x[:, idx[:, 0], idx[:, 1]]
z = torch.transpose(z, 0, 1)
1 Like