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.