Solved by user_123454321 in post #2 
              
              
                I think you need to replace 
nn.Linear(344*11*11, 1376)

with 
nn.Linear(344*10*10, 1376)

Note tha t in max pool the default is to have floored size (that is 21 instead of 21.5 and then 10 instead of 10.5). For example 
>>> import torch
>>> layer = torch.nn.MaxPool2d(2, 2)
>>> a = torch.randn(1, 1,…