Solved by ptrblck in post #2 
              
              
                Most likely you’ll get the tensor shapes for the tensors where the mismatch occurred. 
In your case, it’s: 

RuntimeError: size mismatch, m1: [1 x 26624], m2: [1664 x 1] at … 

If you change the in_features of your linear layer to 26624, your code will run successfully with 320x320 shaped inputs. 
A…