Maybe it is a bit unfortunate that nonzero has the dimensions in that order… 
You can make the split explicit and have something that is close to NumPy’s nonzero indexing:
a = torch.threshold(torch.randn(5, 5, 5), 1, 0)
idx = torch.nonzero(a).split(1, dim=1)
a[idx]

Best regards
Thomas