Keras ConvLSTM InvalidArgumentError - keras

Question

I'm new enough to TensorFlow and Keras that I might be missing something obvious, but this is driving me nuts. I inherited an app that trains a custom convolutional LSTM, and I just spent the last two months or so straightening out some really atrocious data wrangling, only to discover I can't get the model to train properly.
Here's the model definition (with hard values substituted for the variables in my actual code):
inputs = layers.Input(shape = (2, 23, 23, 10))
outputs = layers.ConvLSTM2D(filters = 32,
kernel_size = (5,5),
padding = "same",
return_sequences = True,
stateful = False,
activation = "relu")(inputs)
outputs = layers.BatchNormalization()(outputs)
outputs = layers.ConvLSTM2D(filters = 32,
kernel_size = (3,3),
padding = "same",
return_sequences = True,
stateful = False,
activation = "relu")(outputs)
outputs = layers.ConvLSTM2D(filters = 32,
kernel_size = (3,3),
padding = "same",
return_sequences = True,
stateful = False,
activation = "relu")(outputs)
outputs = layers.Conv3D(filters = 1,
kernel_size = (3, 3, 3),
padding = "same",
activation = "sigmoid")(outputs)
If something looks squirrely there, let me know--I didn't create this model (though I did change the input grid from 100 x 100 to 23 x 23, if that makes a difference). The idea is to predict the intensity of a particular weather phenomenon (that's a single number for each grid point at each time).
Here's the model summary produced after the model is defined:
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) [(None, 2, 23, 23, 10)] 0
conv_lstm2d (ConvLSTM2D) (None, 2, 23, 23, 32) 134528
batch_normalization (BatchN (None, 2, 23, 23, 32) 128
ormalization)
conv_lstm2d_1 (ConvLSTM2D) (None, 2, 23, 23, 32) 73856
conv_lstm2d_2 (ConvLSTM2D) (None, 2, 23, 23, 32) 73856
conv3d (Conv3D) (None, 2, 23, 23, 1) 865
=================================================================
Total params: 283,233
Trainable params: 283,169
Non-trainable params: 64
_________________________________________________________________
The model is fit using data from a custom Sequence sub-class, which produces X input of shape ([batch size], 2, 23, 23, 10) and Y input of shape ([batch size], 2, 23, 23, 1). The batch size is usually 8, but because of the way the data is lazily loaded, the last batch in a particular block of files may be smaller, which is why I don't specify a batch size in model definition. For the record, the original coders had a constant batch size, though, as with mine, it wasn't specified in the model definition.
When I try to fit the model, I get a crash pretty quickly, with this traceback:
Traceback (most recent call last):
File "C:/code/Python/edapts/ConvLSTM2D.py", line 175, in <module>
history = model.fit(training_data,
File "C:\Python\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "C:\Python\lib\site-packages\tensorflow\python\eager\execute.py", line 54, in quick_execute
tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:
Detected at node 'gradient_tape/model/conv_lstm2d/transpose_1/transpose' defined at (most recent call last):
File "C:/code/Python/edapts/ConvLSTM2D.py", line 175, in <module>
history = model.fit(training_data,
File "C:\Python\lib\site-packages\keras\utils\traceback_utils.py", line 64, in error_handler
return fn(*args, **kwargs)
File "C:\Python\lib\site-packages\keras\engine\training.py", line 1409, in fit
tmp_logs = self.train_function(iterator)
File "C:\Python\lib\site-packages\keras\engine\training.py", line 1051, in train_function
return step_function(self, iterator)
File "C:\Python\lib\site-packages\keras\engine\training.py", line 1040, in step_function
outputs = model.distribute_strategy.run(run_step, args=(data,))
File "C:\Python\lib\site-packages\keras\engine\training.py", line 1030, in run_step
outputs = model.train_step(data)
File "C:\Python\lib\site-packages\keras\engine\training.py", line 893, in train_step
self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
File "C:\Python\lib\site-packages\keras\optimizers\optimizer_v2\optimizer_v2.py", line 537, in minimize
grads_and_vars = self._compute_gradients(
File "C:\Python\lib\site-packages\keras\optimizers\optimizer_v2\optimizer_v2.py", line 590, in _compute_gradients
grads_and_vars = self._get_gradients(tape, loss, var_list, grad_loss)
File "C:\Python\lib\site-packages\keras\optimizers\optimizer_v2\optimizer_v2.py", line 471, in _get_gradients
grads = tape.gradient(loss, var_list, grad_loss)
Node: 'gradient_tape/model/conv_lstm2d/transpose_1/transpose'
transpose expects a vector of size 4. But input(1) is a vector of size 5
[[{{node gradient_tape/model/conv_lstm2d/transpose_1/transpose}}]] [Op:__inference_train_function_9880]
I'm so lost. The only thing I can think of is that the original coders made the X and Y features the same, whereas I'm only trying to predict 1 of the 10 (there's not actually a reason to predict the others--not sure why they were trying to). If that's the problem, how do I redefine the model to take account of the different output shape?
EDIT: Well, that's interesting. I just compared the model definition from the original team with the model definition used by the co-worker I inherited the code from. Said co-worker seems to have inserted an extra layer in the outputs, duplicating the ConvLSTM2D with 32 filters and kernel size (3,3). When I remove one of those duplicates, everything runs just fine....
Great, so it's "fixed". But is there someone who can explain why it wasn't working in the first place? My level of understanding the issue at this point is to cross myself and throw salt over my shoulder.
EDIT #2: Does the problem result from having such a small grid (23 x 23)? So either a bigger grid or smaller kernels would solve the problem without deleting a layer? That seems intuitively likely, but I'd like to know the match to calculate how the outputs from each layer match up with the definition of the next layer.