I’m not 100% sure I get how to solve that. I removed the softmax and changed the x = x.view(-1, self._to_linear) to x = x.view(x.size(0), -1) as you recommended, and I do get a shape mismatch, but I don’t understand which layer is having the issue, and I don’t get why the batch_size of 10 once again finds its way into the error.
Here is the error message by the way:
RuntimeError: size mismatch, m1: [10 x 100352], m2: [512 x 512] at C:\w\1\s\tmp_conda_3.7_104508\conda\conda-bld\pytorch_1572950778684\work\aten\src\TH/generic/THTensorMath.cpp:197

Is the problem cause in the forward or in the layer itself?
Here is a link to my updated code on github after all the small changes I’ve done with you guys’ help if the format is getting difficult to follow:
Code on Github 10