title: AssertionError: No inf checks were recorded for this optimizer in Pytorch's AutomaticMixedPrecision
body: <p>I'm using AutomaticMixedPrecision feature of PyTorch to train a network with smaller footprint and precision.<br />
  At a certain point some embeddings from the network have NaNs in their tensors, so I'd like to replace those with 0s in order to perform online hard negative samples mining.</p>
  <p>However, after replacing the NaNs in the tensor like this:</p>
  <pre><code>tensor[torch.isnan(tensor)] = 0
  </code></pre>
  <p>I get the following error while doing the next scaler ste (scaler.step(optimizer):</p>
  <pre><code>    assert len(optimizer_state[&quot;found_inf_per_device&quot;]) &gt; 0, &quot;No inf checks were recorded for this optimizer.&quot;
  AssertionError: No inf checks were recorded for this optimizer.
  </code></pre>
  <p>What's the correct way to zero out NaNs while getting rid of this error?</p>