Solved by ptrblck in post #2 
              
              
                Since you are wrapping all submodules into an nn.Sequential module, you are losing the flatten operation which is performed in the forward method. 
Similar to this issue you can add a custom Flatten module and it should work: 
class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, …