I think that the error is originated by the wrong use of the .batch method (I don't know why it doesn't make the image shape = (128, 32, 32, 3) ) or by the wrong definition of the estimator.
In the original code (taken from this repository 
CodePudding user response:
The solution to the wrong shape of the dataset (being (?, 32, 32, 3) instead of (128, 32, 32, 3) ) is to add drop_remainder=True to the args of the batch() method.
This seems to solve the problem.
