Ошибка несовместимых фигур на fit_generator

Question

Мне нужно обучить модель VAE, используя метод fit_generator, но когда размер партии не равен 1, я получаю ошибку «Несовместимые формы [высота * ширина * размер партии] и [размер партии]».Я понимаю, почему это происходит, но я не знаю, где это исправить.

Я взял код из книги Франсуа Шолле ( ссылка на GitHub ), в примере подгонкаметод использовался, но с помощью fit_generator код не работал без изменений, я написал отдельную функцию потерь.

import keras
from keras import layers
from keras import backend as K
from keras.models import Model
import numpy as np

img_shape = (180, 220, 1)
latent_dim = 3

input_img = keras.Input(shape=img_shape)

x = layers.Conv2D(32, 3,
              padding='same', activation='relu')(input_img)
x = layers.Conv2D(64, 3,
              padding='same', activation='relu',
              strides=(2, 2))(x)
x = layers.Conv2D(64, 3,
              padding='same', activation='relu')(x)
x = layers.Conv2D(64, 3,
              padding='same', activation='relu')(x)
x = layers.Conv2D(64, 3,
              padding='same', activation='relu',
              strides=(2, 2))(x)
x = layers.Conv2D(64, 3,
              padding='same', activation='relu')(x)
shape_before_flattening = K.int_shape(x)

x = layers.Flatten()(x)
x = layers.Dense(32, activation='relu')(x)

z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)

def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim),
                          mean=0., stddev=1.)
    return z_mean + K.exp(z_log_var) * epsilon

z = layers.Lambda(sampling)([z_mean, z_log_var])

decoder_input = layers.Input(K.int_shape(z)[1:])

x = layers.Dense(np.prod(shape_before_flattening[1:]),
             activation='relu')(decoder_input)

x = layers.Reshape(shape_before_flattening[1:])(x)

x = layers.Conv2DTranspose(32, 3,
                       padding='same', activation='relu',
                       strides=(2, 2))(x)
x = layers.Conv2DTranspose(32, 3,
                       padding='same', activation='relu',
                       strides=(2, 2))(x)
x = layers.Conv2D(1, 3,
              padding='same', activation='sigmoid')(x)

decoder = Model(decoder_input, x)

z_decoded = decoder(z)

class VariationalLayer(keras.layers.Layer):

    def vae_loss(self, x, z_decoded):
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        return K.mean(xent_loss + kl_loss)

    def call(self, inputs):
        x = inputs[0]
        z_decoded = inputs[1]
        loss = self.vae_loss(x, z_decoded)
        self.add_loss(loss, inputs=inputs)
        return x

y = VariationalLayer()([input_img, z_decoded])

 def vae_loss(x, z_decoded):
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        return K.mean(xent_loss + kl_loss)

vae = Model(input_img, y)
vae.compile(optimizer='rmsprop', loss=vae_loss)

from keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
                    '/content/img',
                    target_size=(180, 220),
                    color_mode="grayscale",
                    batch_size=1,
                    class_mode='categorical')
vae.fit_generator(train_generator,
              steps_per_epoch=22511,
              epochs=1)

Конкретная ошибка в приведенном выше коде:

InvalidArgumentError: Incompatibleформы: [356400] против [9]