error with a model with some blanchs in input layer

[oscesteros]

hello contributors:

I've tried to use Lucid with a customized model using the .save function. This model has an input layer that is split into three branches (red, green, and blue are processed in separated branches and concatenate at the end). I have used a lambda function and a reshape:

Input = tf.keras.Input(shape=(224,224,3))
batch=tf.shape(Input)
R= tf.keras.layers.Lambda(lambda x: x[:,:,:,0],)(Input)
R= tf.reshape( R, [batch[0], 224,224,1])

I have obtained this error which the .render_vis in reshape operation:

InvalidArgumentError: 2 root error(s) found.
(0) Invalid argument: Input to reshape is a tensor with 21025 values, but the requested shape has 50176
[[{{node import/Reshape}}]]
[[Mean/_29]]
(1) Invalid argument: Input to reshape is a tensor with 21025 values, but the requested shape has 50176
[[{{node import/Reshape}}]]
0 successful operations.
0 derived errors ignored.

the image has 224x224, and 50176 are values for a tensor with a channel (224x224=50176). 21025 is a 145x145 matrix.

thanks

Oscar