Hi all,
I’m using torch.where to avoid backpropagating through branches that might yield NaN values. Specifically, I want to exponentiate a number if it’s nonnegative and do some other stuff otherwise (as exponentiating a negative number may yield imaginary numbers).
However, when backpropagating from the loss, the resulting gradient is still NaN, even though the loss is the desired one. Here’s a minimal working example. One can use a nullcontext instead of the autograd anomaly detection to check that a.grad is NaN.
a = torch.as_tensor(-0.1).requires_grad_()
with torch.autograd.detect_anomaly():
r = torch.where(a < 0, a * 10, a ** 1.8)
r.backward()
In [77]: r
Out[77]: tensor(-1., grad_fn=<SWhereBackward>)
Any ideas on how to circumvent this? Is this considered a bug?
Thanks in advance.