InvalidArgumentError: Incompatível formas: [3,256,256,2] vs. [3,150,150,2]

Eu quero construir um modelo para colorir a imagem, tentei executar este código, mas eu enfrentei Incompatível formas: [3,256,256,2] vs. [3,150,150,2] erro.

#CNN model

from keras.layers import Conv2D, Conv2DTranspose, UpSampling2D
from keras.layers import Activation, Dense, Dropout, Flatten, InputLayer
from tensorflow.keras.layers import (
    BatchNormalization, SeparableConv2D, MaxPooling2D, Activation, Flatten, Dropout, Dense
)
from keras.callbacks import TensorBoard, ModelCheckpoint
from keras.models import Sequential

model = Sequential()

#Input Layer
model.add(Conv2D(64, (3, 3), input_shape=(256, 256, 1), activation='relu', padding='same'))

#Hidden Layers
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))



#Compiling the CNN
model.compile(optimizer='rmsprop', loss='mse', metrics = ['accuracy'])
#model.compile(optimizer='rmsprop', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

o erro ocorre ao executar o código para ajustar o modelo

# Image transformer
datagen = ImageDataGenerator(
        shear_range=0.2,
        zoom_range=0.2,
        rotation_range=20,
        horizontal_flip=True)

import matplotlib.pyplot as plt

# Image transformer
datagen = ImageDataGenerator(
        shear_range=0.2,
        zoom_range=0.2,
        rotation_range=20,
        horizontal_flip=True)

# Generate training data
batch_size = 10
def image_a_b_gen(batch_size):
    for batch in datagen.flow(Xtrain, batch_size=batch_size ):
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        Y_batch = lab_batch[:,:,:,1:] / 128
        yield (X_batch.reshape(X_batch.shape+(1,)), Y_batch)


        # Train model      
tensorboard = TensorBoard(log_dir="/output/beta_run")
trainedmodel = model.fit(image_a_b_gen(batch_size), callbacks=[tensorboard],epochs=100, steps_per_epoch=30)

Mensagem de erro:

 InvalidArgumentError                      Traceback (most recent call last)
    <ipython-input-112-7a987e785f95> in <module>
         29         # Train model
         30 tensorboard = TensorBoard(log_dir="/output/beta_run")
    ---> 31 trainedmodel = model.fit(image_a_b_gen(batch_size), callbacks=[tensorboard],epochs=100, steps_per_epoch=30)
         32 
         33 
    
    ~\anaconda3\lib\site-packages\keras\utils\traceback_utils.py in error_handler(*args, **kwargs)
         65     except Exception as e:  # pylint: disable=broad-except
         66       filtered_tb = _process_traceback_frames(e.__traceback__)
    ---> 67       raise e.with_traceback(filtered_tb) from None
         68     finally:
         69       del filtered_tb
    
    ~\anaconda3\lib\site-packages\tensorflow\python\eager\execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
         56   try:
         57     ctx.ensure_initialized()
    ---> 58     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
         59                                         inputs, attrs, num_outputs)
         60   except core._NotOkStatusException as e:
    
    InvalidArgumentError:  Incompatible shapes: [3,256,256,2] vs. [3,150,150,2]
         [[node gradient_tape/mean_squared_error/BroadcastGradientArgs
     (defined at C:\Users\HudaA\anaconda3\lib\site-packages\keras\optimizer_v2\optimizer_v2.py:464)
    ]] [Op:__inference_train_function_33345]

Tentei imprimir a resumir de camadas, mas não consegui encontrar o problema