Solved by ptrblck in post #2 
              
              
                The view operation is wrong and will change the batch size. 
Use out = out.view(out.size(0), -1) instead and set the in_features of self.fc1 to 128*128*64.