I am trying to reduce a tensor down to only the non-zero elements. If I am reading the docs correctly I should be able to do the following
# y_true has torch.Size([1, 6, 480, 640])
idx = torch.nonzero(torch.any(torch.greater(y_true, 0.0), dim=1), as_tuple=True)
y_true = y_true[idx] 

However, when I do, I get this error
IndexError: index 324 is out of bounds for dimension 1 with size 6

What have I missed?