Hi @Cactus, try explicitly setting the input shape of the Flatten layer like Flatten(input_shape=(28, 28)). Right now your layer is inferring the first dimension from the default batch_size that the fit method has but it should be agnostic from the batch size used.