Solved by ptrblck in post #2 
              
              
                Calling model.half() can easily create under/overflows, which could explain the NaN values in the forward/backward pass, which is why we recommend using torch.cuda.amp for mixed-precision training. 
Are you seeing the first cudnn error using amp? 
If so, could you post a minimal, executable code sni…