cartpole_reinforce.ipynb 에러

코랩에서 cartpole_reinforce.ipynb 실행시킬때 아래와 같은 에러 발생합니다.

ValueError: Layer "my_model_1" expects 3 input(s), but it received 1 input tensors. Inputs received: [<tf.Tensor 'IteratorGetNext:0' shape=(None, 1, 4) dtype=float32>]

에러나는 소스 위치는 아래 화면에서 "y_pred = self(states, training=True)" 부분에서

에러가 발생됩니다.

    class MyModel(tf.keras.Model):
        def train_step(self, data):
            in_datas, out_actions = data
            states, action_matrix, rewards = in_datas[0], in_datas[1], in_datas[2]

            with tf.GradientTape() as tape:
                y_pred = self(states, training=True)
                action_probs = K.sum(action_matrix*y_pred, axis=-1)
                loss = -K.log(action_probs)*rewards

            trainable_vars = self.trainable_variables
            gradients = tape.gradient(loss, trainable_vars)
            self.optimizer.apply_gradients(zip(gradients, trainable_vars))

 tensorflow 버젼 : 2.9.2 이며, 모델에 3개의 값을 주어야 하는데, 1개만 받아서 그런것 같습니다.

해결 방법 알려 주시기 바랍니다.

감사합니다.