title: Invalid shape (64, 125, 125, 3) for image data after using ImageDataGenerator
body: <p>So I'm trying to display image that already augmented. But got <code>Invalid shape (64, 125, 125, 3) for image data</code> error. This is my code:</p>
  <pre><code>from tensorflow.keras.preprocessing.image import ImageDataGenerator
  
  train_datagen = ImageDataGenerator(rescale=1./255,
                                      zoom_range=0.1,
                                      rotation_range=25,
                                      width_shift_range=0.1,
                                      height_shift_range=0.1,
                                      shear_range=0.1, horizontal_flip=True,
                                      fill_mode='nearest')
  
  val_datagen = ImageDataGenerator(rescale=1./255)
  
  # build image augmentation generators
  train_generator = train_datagen.flow(train_data, train_labels_enc, batch_size=64, shuffle=True)
  val_generator = val_datagen.flow(val_data, val_labels_enc, batch_size=64, shuffle=False)
  
  
  from matplotlib import pyplot
  for i in range(9):
      # define subplot
      pyplot.subplot(330 + 1 + i)
      # generate batch of images
      batch = train_generator.next()
      # convert to unsigned integers for viewing
      image = batch[0].astype('uint8')
      # plot raw pixel data
      pyplot.imshow(image)
  
  pyplot.show()
  
  </code></pre>
  <p>And the error pointed to <code>pyplot.imshow(image)</code>. Any help would be appreciated</p>