Here is a numpy example to use
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