The TypeError in BroadcastTo.call() was caused by the Masking layer (applied to jet_masked = keras.layers.Masking(mask_value=0.0)(jet_input)). This layer created a boolean mask (shape (None, 3, 2)) that interfered with downstream layers like LSTM. To fix this, remove the Masking layer and instead pass a custom jet_mask tensor directly to your LSTM's mask argument. This approach prevents automatic masking while still allowing jet_input to be concatenated with other inputs. Please refer to this gist.