tensorflow - TF / Keras エラー: InputLayer は Checkpointable ではありません

Question

単純な猫と犬のデータセットを使用して、Google Colab で新しく追加された TPU サポートを試しています。

簡単な CNN を作成した後、モデルを TPU にエクスポートしようとしました。しかし、それはエラーで失敗しました

TypeError: Checkpointable._track_checkpointable() passed type <class 'keras.engine.topology.InputLayer'>, not a Checkpointable.

Colabで書いたコードは次のとおりです。

model = models.Sequential()
model.add(layers.Conv2D(32, (3,3), activation='relu', input_shape=(150, 150, 3)))
model.add(layers.MaxPooling2D((2,2)))
model.add(layers.Conv2D(64, (3,3), activation='relu'))
model.add(layers.MaxPooling2D((2,2)))
model.add(layers.Flatten())
model.add(layers.Dropout(0.5))
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))
model.summary()

train_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(train_dir, target_size=(150,150), batch_size=20, class_mode='binary')

tpu_model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=tf.contrib.tpu.TPUDistributionStrategy(tf.contrib.cluster_resolver.TPUClusterResolver(tpu="grpc://" + os.environ['COLAB_TPU_ADDR'])))

私の推測では、 で何か間違ったことをしているのですtrain_generator。しかし、それが何であるかはわかりません。どんな助けでも大歓迎です。