Solved by jpeg729 in post #2 
              
              
                You should be able to make batches of images of shape (n_images, 1, 50, 50) and train with those. 
My guess is that the code you use to flatten the cnn output is incorrect, but just happens to work when sending images 1 by 1. My guess is that your code turns the 4d output of the cnn into a 1d tensor…