Output of ResNet-18 is NaN with AMP

username: D0miH
cooked: <p>Hello,<br>
I want to use AMP on a ResNet-18 which was trained without AMP (plain Float32) on CIFAR-10.<br>
However, when I wrap the forward pass of the model in a <code>torch.cuda.amp.autocast()</code> block, the output of the network gets <code>nan</code>. When I deactivate AMP with <code>torch.cuda.amp.autocast(enabled=False)</code> I get the expected output values.<br>
Below I attached a minimal example that reproduces my problem.</p>
<p>Does anyone have an idea why the output is always <code>nan</code>?</p>
<pre><code class="lang-python">import torch

from models.cifar10_models.resnet import ResNet18

resnet = ResNet18(num_classes=10)
resnet.load_state_dict(torch.load('./pretrained_models/resnet18_cifar10.pth'))

inputs = torch.randn((124, 3, 32, 32), requires_grad=True).cuda()

resnet.cuda()
resnet.eval()
for i in range(500):
    with torch.cuda.amp.autocast():
        outputs = resnet(inputs)

print(outputs)
</code></pre>
<p>Edit:<br>
I have noticed that reproducing the problem is difficult when I am loading a local model. Therefore, here another example with a pretrained ResNet-18 on ImageNet. Running the code below also produces <code>nan</code> as output:</p>
<pre><code class="lang-python">import torch
from torchvision.models.resnet import resnet18

resnet = resnet18(pretrained=True)
resnet.fc = torch.nn.Linear(512, 10)

inputs = torch.randn((124, 3, 32, 32), requires_grad=False).cuda()

resnet.cuda()
resnet.eval()
with torch.cuda.amp.autocast(enabled=True):
    outputs = resnet(inputs)

print(outputs)
</code></pre>

username: RaLo4
cooked: <p>Unfortunately I cannot reproduce the issue.<br>
No matter how often I try I never get <code>nan</code><br>
This might be a version problem.</p>
<p>It tried it on:<br>
CUDA 10.2 / Pytorch1.8.1<br>
as well as on:<br>
CUDA 10.2 / Pytorch1.9.0</p>
<p>What are you using?</p>

username: D0miH
cooked: <p>Oh, you are right. If I am using the <code>pytorch/pytorch:1.9.0-cuda10.2-cudnn7-devel</code> docker image to run the example the outputs are fine.</p>
<p>In my environment I have installe <code>torch==1.9.0</code>, <code>torchvision=0.10.0</code> and <code>cudotoolkit=11.1</code> using conda. Seems like the problem with the NaN-Output only happens when using Cuda 11.1. When I install <code>cudatoolkit=10.2</code> everything works fine.</p>
<p>Thanks for your help! Very much appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=10" title=":smiley:" class="emoji" alt=":smiley:"> <img src="https://discuss.pytorch.org/images/emoji/apple/+1.png?v=10" title=":+1:" class="emoji" alt=":+1:"></p>

username: ptrblck
cooked: <p>Which GPU are you using to create the NaN outputs in PyTorch <code>1.9.0</code>?</p>

username: D0miH
cooked: <aside class="quote no-group" data-username="D0miH" data-post="1" data-topic="128328">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/d0mih/48/12990_2.png" class="avatar"> D0miH:</div>
<blockquote>
<pre><code class="lang-auto">import torch
from torchvision.models.resnet import resnet18

resnet = resnet18(pretrained=True)
resnet.fc = torch.nn.Linear(512, 10)

inputs = torch.randn((124, 3, 32, 32), requires_grad=False).cuda()

resnet.cuda()
resnet.eval()
with torch.cuda.amp.autocast(enabled=True):
    outputs = resnet(inputs)

print(outputs)
</code></pre>
</blockquote>
</aside>
<p>I am using a GTX 1650.</p>