Thanks for the code snippet. I cannot reproduce the NaN outputs for the amp model and get:
7.946749687194824
7.837947368621826
8.126112937927246
5.472558498382568
3.0419869422912598
1.3717389106750488
0.4931238293647766
0.21829186379909515
0.09434405714273453
0.0426550917327404
0.021307891234755516
0.013564658351242542
0.010585908778011799
0.00768495025113225
0.005186968017369509
0.0036425101570785046
0.003632057225331664
0.0027618384920060635
0.0024199967738240957
0.00201830524019897
0.0015914703253656626
0.0011888183653354645
0.0011378041235730052
0.000954880437348038
0.0007914779707789421
0.000882966909557581
0.0007371000247076154
0.0007495403406210244
0.0006300751701928675

Could you post the PyTorch, torchvision, CUDA, and cudnn versions you are using and how you’ve installed PyTorch?