The torch.nonzero() operation should probably be replaced with torch.where, as you are using np.where in your original code.
Besides that it looks alright, but I would recommend to verify it with some dummy tensors.
Just create random tensors in the expected shape, and compare the outputs of the new PyTorch method to the original numpy method (by passing them as tensor.numpy() into the function).