Solved by ptrblck in post #2 
              
              
                If you wrap all modules in an nn.Sequential module, you are not using the original forward method anymore, and thus are missing this reshape operation. 
You could add a custom Flatten layer or write your own forward method using a custom nn.Module: 
class Flatten(nn.Module):
    def __init__(self):
…