callback 문제로 질문 드립니다.

import urllib.request
import zipfile
import tensorflow.keras as keras

url = 'https://storage.googleapis.com/download.tensorflow.org/data/rps.zip'
urllib.request.urlretrieve(url, 'rps.zip')
local_zip = 'rps.zip'
zip_ref = zipfile.ZipFile(local_zip, 'r')
zip_ref.extractall('tmp/')
zip_ref.close()


def image_generator():
    TRAIN_DIR = 'tmp/rps'
    training_datagen = keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255,
                                                                    rotation_range=40,
                                                                    width_shift_range=0.2,
                                                                    height_shift_range=0.2,
                                                                    shear_range=0.2,
                                                                    zoom_range=0.2,
                                                                    horizontal_flip=True,
                                                                    fill_mode='nearest',
                                                                    validation_split=0.2)
    training_generator = training_datagen.flow_from_directory(TRAIN_DIR,
                                                              target_size=(128, 128),
                                                              batch_size=150,
                                                              class_mode='categorical',
                                                              subset='training')

    validation_generator = training_datagen.flow_from_directory(TRAIN_DIR,
                                                                target_size=(128, 128),
                                                                batch_size=150,
                                                                class_mode='categorical',
                                                                subset='validation')

    return training_generator, validation_generator


model = keras.models.Sequential([
        keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(150, 150, 3)),
        keras.layers.MaxPooling2D(2, 2),
        keras.layers.Conv2D(32, (3, 3), activation='relu'),
        keras.layers.MaxPooling2D(2, 2),
        keras.layers.Conv2D(16, (3, 3), activation='relu'),
        keras.layers.MaxPooling2D(2, 2),
        keras.layers.Flatten(),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(512, activation='relu'),
        keras.layers.Dense(3, activation='softmax')
    ])

model.compile(optimizer=keras.optimizers.Adam(),
              loss=keras.losses.sparse_categorical_crossentropy,
              metrics='acc')


checkpoint_path = 'models/my_checkpoint.ckpt'
checkpoint = keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                             save_best_only=True,
                                             save_weights_only=True,
                                             monitor='val_loss',
                                             verbose=1)


train, val = image_generator()

model.fit(train,
          epochs=10,
          validation_data=(val),
          steps_per_epoch=len(train),
          validation_steps=len(val),
          callbacks=[checkpoint])

model.load_weights(checkpoint_path)

tensorflow.python.framework.errors_impl.InvalidArgumentError:  Input to reshape is a tensor with 470400 values, but the requested shape requires a multiple of 4624

[[node sequential/flatten/Reshape (defined at /Users/resot/PycharmProjects/TensorProject/imageclassification_cnn.py:72) ]] [Op:__inference_train_function_874]

Function call stack:

train_function

실습자료를 하던 중

callback에서 걸려서 이런 오류가 뜹니다.

폴더와 파일생성은 모두 되었습니다.

flatten에서 reshape를 해줘야 하는건가요? 

강의자료와 코드를 비교해보아도 알 수 없어 질문을 드립니다.