Here is a numpy example to use nonzero.
    a = np.random.randint(6,size=(4,5,3))
    idx = np.nonzero(a)
    a[idx] = 0

This is PyTorch
    a = torch.randint(6,size=(4,5,3))
    idx = torch.nonzero(a)  # idx = a.nonzero()
    # a[idx] = 0  # this throws an error because `nonzero` cannot be used as index
    a[(idx[:, 0], idx[:, 1], idx[:, 2])] = 0  # this is a little counter-intuitive, can we instead accept multi-dim tensors as index?


I actually like torch.nonzero's single tensor better than the tuple that numpy returns, it is more elegant. But I do not know of a better way to use it as index