tensorflow: Exception encountered when calling layer of sub-classe model

Question

I try to build a multivariate timeseries forcast lstm using auto-regressive of one of my features in the past to forcast it (all other feature are available in futur) I use this exemple for the main idea :https://www.tensorflow.org/tutorials/structured_data/time_series?hl=fr#advanced_autoregressive_model

my model code is like this:

def custom_model(self):

        class FeedBack(tf.keras.Model):
            def __init__(self, units, in_steps, out_steps):
                super().__init__()
                self.in_steps = in_steps
                self.out_steps = out_steps
                self.units = units
                self.lstm_cell = tf.keras.layers.LSTMCell(units[0])
                self.drop1 = tf.keras.layers.Dropout(.2)
                self.drop2 = tf.keras.layers.Dropout(.2)
                # Also wrap the LSTMCell in an RNN to simplify the `warmup` method.
                self.lstm_rnn = tf.keras.layers.RNN(self.lstm_cell, return_sequences=False, return_state=True)
                self.dense = tf.keras.layers.Dense(30)
                self.dense_out = tf.keras.layers.Dense(1)
            
            def warmup(self, seq_feature, seq_label, window):
            # use a x timestep sequence to innitialize the model

                # need to shift one feature (last one) of one step
                ft_indx = list(range(1, window)) 
                lbl_indx = list(range(0, window-1))  

                window_inputs = tf.concat([tf.gather(seq_feature, ft_indx, axis=1), 
                                                tf.expand_dims(tf.gather(seq_label, lbl_indx, axis=1), 2)], 2)
                # inputs.shape => (batch, time, features)
                # x.shape => (batch, lstm_units)
                x, *state = self.lstm_rnn(window_inputs)
                x = self.drop1(x)

                # predictions.shape => (batch, features)
                x = self.dense(x)
                x = self.drop2(x)
                prediction = self.dense_out(x)
                return prediction, state


            def call(self, inputs, training=None):
                # Use a TensorArray to capture dynamically unrolled outputs.
                predictions = []
                timeseries_input=inputs[0]

                attributs_inputs = inputs[1]
                inp_shape = timeseries_input.get_shape().as_list()
                # inp_shape = [None, timesteps, features]

                # select features from dataset
                seq_feature = tf.gather(timeseries_input, list(range(0, inp_shape[2]-1)), axis=2)
                # select labels use as features with a shift of one timestep
                seq_label = tf.gather(timeseries_input, inp_shape[2], axis=2)   

                # Initialize the LSTM state.
                prediction, state = self.warmup(seq_feature, seq_label, self.in_steps)
                # Insert the first prediction.
                predictions.append(prediction)
                idx=0

                # Run the rest of the prediction steps.
                for n in range(self.in_steps, inp_shape[1]):
                    # forcast is done one timestep at time using last forcast as the new seq_label
                    # concat feature and last forcast
                    window_features = tf.concat([tf.expand_dims(tf.gather(seq_feature, n, axis=1),1), 
                                                 tf.expand_dims(predictions[idx],1)], 2)
            

                    # Execute one lstm step.
                    x, *state = self.lstm_rnn(window_features, initial_state=state,
                                    training=training)
                    x = self.drop1(x)

                    # predictions.shape => (batch, features)
                    x = self.dense(x)
                    x = self.drop2(x)
                    prediction = self.dense_out(x)
                    # Add new the prediction.
                    predictions.append(prediction)
                    # update indexer
                    idx = idx+1
                # predictions.shape => (time, batch, features)
                predictions = tf.stack(predictions)
                # predictions.shape => (batch, time, features)
                predictions = tf.transpose(predictions, [1, 0, 2])
                return predictions
            
        # Model
        self.NN = FeedBack(units=self.rnn_layers_size, in_steps=self.nb_timestep_in, out_steps=self.nb_timestep_out)
        self.NN.build(input_shape=[(None,self.nb_timestep_in + self.nb_timestep_out, self.nb_features-self.nb_attributs),
                                    (None, self.nb_timestep_out, self.nb_attributs)])

my dataset input#1 have 4 features, input#2 have 15 but not used in this model exemple. The output are sequence of 10 but not fixed shape are:

_type_spec: DatasetSpec(({'input1': TensorSpec(shape=(None, None, 4), dtype=tf.float32, name=None), 'input2': TensorSpec(shape=(None, 15), dtype=tf.float32, name=None)}, TensorSpec(shape=(None, None), dtype=tf.float32, name=None)), TensorShape([]))

get this error:

Exception has occurred: RuntimeError
in user code:

    File "/home/jo/miniconda3/envs/gpu/lib/python3.10/site-packages/keras/src/engine/training.py", line 1377, in train_function  *
        return step_function(self, iterator)
    File "/home/jo/miniconda3/envs/gpu/lib/python3.10/site-packages/keras/src/engine/training.py", line 1360, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/jo/miniconda3/envs/gpu/lib/python3.10/site-packages/keras/src/engine/training.py", line 1349, in run_step  **
        outputs = model.train_step(data)
    File "/home/jo/miniconda3/envs/gpu/lib/python3.10/site-packages/keras/src/engine/training.py", line 1126, in train_step
        y_pred = self(x, training=True)
    File "/home/jo/miniconda3/envs/gpu/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_fileccrhp_wh.py", line 11, in tf__call
        timeseries_input = ag__.ld(inputs)[0]

    RuntimeError: Exception encountered when calling layer 'feed_back' (type FeedBack).
    
    KeyError: 0
    
    Call arguments received by layer 'feed_back' (type FeedBack):
      • inputs={'input1': 'tf.Tensor(shape=(None, None, 4), dtype=float32)', 'input2': 'tf.Tensor(shape=(None, 15), dtype=float32)'}
      • training=True
  File "/tmp/__autograph_generated_fileccrhp_wh.py", line 11, in tf__call
    timeseries_input = ag__.ld(inputs)[0]
KeyError: 0