`torch.mean` return nan when enable `torch.cuda.amp.autocast`

username: pytorching
cooked: <p>My code works when disable <code>torch.cuda.amp.autocast</code>.<br>
But when I enable <code>torch.cuda.amp.autocast</code>,  I found that the normalization step overflows (float16)<br>
<code>x = x - x.mean()</code>, where <code>x.shape = [5665, 48]</code> and <code>x.mean()</code> returns <code>nan</code>.<br>
Is there any way to solve this problem?<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee03c1df9945e18a327f245b9d33ad01859a903c.png" alt="image" data-base62-sha1="xXzGEBNk8HS0Gm9MG1vjzxAcNqI" width="629" height="435"></p>

username: ptrblck
cooked: <p>I cannot reproduce the issue by trying to create a tensor in a similar range:</p>
<pre><code class="lang-python">x = torch.randn(5665, 48, device='cuda', dtype=torch.float16) * 100.
print(x)
with torch.cuda.amp.autocast():
    out = x - x.mean()
print(torch.isnan(out).any())
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">tensor([[  68.5000,   57.2188,  -15.8438,  ...,  -59.6562,  -81.4375,
          -58.6562],
        [  48.5938, -110.9375,  112.6250,  ..., -133.3750,   12.9766,
          -29.9062],
        [  -9.2812,  -36.3125,  -98.5000,  ...,  139.5000,  -30.3438,
           59.9062],
        ...,
        [ 114.8750,   89.8125,   23.7500,  ..., -150.6250,  -46.2500,
          -32.7500],
        [  42.2500,  -17.6406,  103.8750,  ...,  -55.9062,   38.0312,
          -62.6875],
        [  53.4688,  187.7500,   44.3125,  ...,  122.9375,   40.6875,
         -123.4375]], device='cuda:0', dtype=torch.float16)
tensor(False, device='cuda:0')
</code></pre>

username: pytorching
cooked: <p>I use <code>torch.save(tensor, 'x.pt')</code> to save my tensor.<br>
Could you analyze the overflow reason for me? I don’t know whether it’s allowed to share links to download files.<br>
My dumped tensor <code>x</code>:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://wetransfer.com/downloads/be87ae199bd48ee2d2ba300714090b3820210827122648/14acce">
  <header class="source">
      <img src="https://prod-cdn.wetransfer.net/packs/media/images/favicon-a34a7465.ico" class="site-icon" width="64" height="64">

      <a href="https://wetransfer.com/downloads/be87ae199bd48ee2d2ba300714090b3820210827122648/14acce" target="_blank" rel="noopener nofollow ugc">wetransfer.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://prod-cdn.wetransfer.net/packs/media/images/wt-facebook-9db47080.png" class="thumbnail onebox-avatar" width="200" height="200">

<h3><a href="https://wetransfer.com/downloads/be87ae199bd48ee2d2ba300714090b3820210827122648/14acce" target="_blank" rel="noopener nofollow ugc">x.pt</a></h3>

  <p>1 file sent via WeTransfer, the simplest way to send your files around the world</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
I think the tensor has <code>dtype=torch.half</code>.</p>

username: ptrblck
cooked: <p>Unfortunately, I cannot download the file (the site claims it’s downloading but nothing happens).<br>
I also cannot reproduce the issue using the max. values in <code>float16</code>:</p>
<pre><code class="lang-python">x = torch.ones(5665, 48, device='cuda', dtype=torch.float16) * 65504
print(x)
with torch.cuda.amp.autocast():
    out = x - x.mean()
print(torch.isnan(out).any())
&gt; tensor(False, device='cuda:0')
</code></pre>

username: pytorching
cooked: <p><code>x.mean()</code> output <code>nan</code>, because it has both negative and positive values, which will accumulate to <code>Inf</code> and <code>-Inf</code> and that’s why there’s <code>nan</code> in the tensor. Otherwise, at worst, there will be <code>Inf</code> instead of <code>nan</code>.<br>
Could you try it again with this link? The link above is fine, I can download it without any problem.<br>
<a href="https://cowtransfer.com/s/593ab3ca8d3946" class="onebox" target="_blank" rel="noopener nofollow ugc">https://cowtransfer.com/s/593ab3ca8d3946</a></p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="pytorching" data-post="5" data-topic="130366">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/85f322/48.png" class="avatar"> pytorching:</div>
<blockquote>
<p><code>x.mean()</code> output <code>nan</code>, because it has both negative and positive values, which will accumulate to <code>Inf</code> and <code>-Inf</code> and that’s why there’s <code>nan</code> in the tensor.</p>
</blockquote>
</aside>
<p>No, that’s not the case as seen in my example, as I’m already using the max. <code>float16</code> values and they are not accumulating to <code>Inf</code> (internally <code>float32</code> will be used so the overflow will not happen).</p>
<aside class="quote no-group" data-username="pytorching" data-post="5" data-topic="130366">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/85f322/48.png" class="avatar"> pytorching:</div>
<blockquote>
<p>Could you try it again with this link?</p>
</blockquote>
</aside>
<p>Thanks for the new link as I was able to download the data.<br>
Your input data already contains <code>Inf</code> values, which will create the <code>NaN</code> output in <code>x.mean()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">print(torch.isinf(x).any())
&gt; tensor(True, device='cuda:0')
</code></pre>