Oh, you are right. If I am using the pytorch/pytorch:1.9.0-cuda10.2-cudnn7-devel docker image to run the example the outputs are fine.
In my environment I have installe torch==1.9.0, torchvision=0.10.0 and cudotoolkit=11.1 using conda. Seems like the problem with the NaN-Output only happens when using Cuda 11.1. When I install cudatoolkit=10.2 everything works fine.
Thanks for your help! Very much appreciated