How to replace infs to avoid nan gradients in PyTorch
I need to compute log(1 + exp(x)) and then use automatic differentiation on it. But for too large x, it outputs inf because of the exponentiation:

&gt;&gt;&gt; x = torch.tensor([0., 1., 100.], requires_grad=True)
&gt;&gt;&gt; x.exp().log1p()
tensor([0.6931, 1.3133,    inf], grad_fn=&lt;Log1PBackward&gt;)


Since log(1 + exp(x)) ≈ x for large x, I thought I could replace the infs with x using torch.where. But when doing this, I still get nan for the gradient of too large values. Do you know why this happens and if there is another way to make it work?

&gt;&gt;&gt; exp = x.exp()
&gt;&gt;&gt; y = x.where(torch.isinf(exp), exp.log1p())  # Replace infs with x
&gt;&gt;&gt; y  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=&lt;SWhereBackward&gt;)
&gt;&gt;&gt; y.sum().backward()  # Automatic differentiation
&gt;&gt;&gt; x.grad  # Why is there a nan and how can I get rid of it?
tensor([0.5000, 0.7311,    nan])


A workaround I've found is to manually implement a Log1PlusExp function with its backward counterpart. Yet it does not explain the bad behavior of torch.where in the question.

&gt;&gt;&gt; class Log1PlusExp(torch.autograd.Function):
...     """Implementation of x ↦ log(1 + exp(x))."""
...     @staticmethod
...     def forward(ctx, x):
...         exp = x.exp()
...         ctx.save_for_backward(x)
...         return x.where(torch.isinf(exp), exp.log1p())
...     @staticmethod
...     def backward(ctx, grad_output):
...         x, = ctx.saved_tensors
...         return grad_output / (1 + (-x).exp())
... 
&gt;&gt;&gt; log_1_plus_exp = Log1PlusExp.apply
&gt;&gt;&gt; x = torch.tensor([0., 1., 100.], requires_grad=True)
&gt;&gt;&gt; log_1_plus_exp(x)  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=&lt;Log1PlusExpBackward&gt;)
&gt;&gt;&gt; log_1_plus_exp(x).sum().backward()
&gt;&gt;&gt; x.grad  # And no nans!
tensor([0.5000, 0.7311, 1.0000])

