Selecting sub part of tensor and the indices of selected elements

1

Assume we want to select all the elements in the tensor which are greater than 5. I know that the following code can do that:

B = A[A > 5]

But I want to know the indices of the selected elements. Is there any way to select and find the indices of the selected elements?

Thanks

2

nonzero() should work:

A = torch.randint(0, 10, (10, 10))
idx = (A > 5).nonzero()