title: How to replace infs to avoid nan gradients in PyTorch
body: <p>I need to compute <code>log(1 + exp(x))</code> and then use automatic differentiation on it. But for too large <code>x</code>, it outputs <code>inf</code> because of the exponentiation:</p>
  
  <pre class="lang-py prettyprint-override"><code>&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;)
  </code></pre>
  
  <p>Since <code>log(1 + exp(x)) ≈ x</code> for large <code>x</code>, I thought I could replace the <code>infs</code> with <code>x</code> using <code>torch.where</code>. But when doing this, I still get <code>nan</code> for the gradient of too large values. Do you know why this happens and if there is another way to make it work?</p>
  
  <pre class="lang-py prettyprint-override"><code>&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])
  </code></pre>