Solved by ptrblck in post #2 
              
              
                Remove out = out.view(out.size(0), -1) from your forward method amd you will get an output shape of [batch, 3, 218, 218]. 
.view works like reshape in numpy. So basically you are flattening your output, which gives [batch, 3*218*218=142572]. 
If you need the same width and height (224) in your outpu…