You can find the derivative for torch.where here and for torch.nonzero here, which is not differentiable as its returning indices.
torch.where
torch.nonzero