Solved by Marcel_Braasch in post #3 
              
              
                No, it has nothing to do with tensorflow, I am not using it. Besides my model shown above being incorrect the problem was that I am using an RTX 3090 and my CUDA version was outdated. For others with this problem, running this fixed it: 
$ /usr/bin/python3 -m pip3 install torch==1.7.0+cu110 torchvis…