TeOptimization loop cancelled with Keras tensorflow


I am running a NN on keras with tensorflow 2.6.0. The NN is used for RL. I am getting a message saying Optimization loop failed: Cancelled operation. My intuition is that the program was not able to carry out the optimization loop with gradient descent. The cost function also does not change at all hence no learning is taking place. Is there any solution for this.

Do you have a very minimal Colab or gist to share with dummy inputs?

Here is some code. The function below is the main function that builds the model.

> def build_model(self):
>         state_input = Input((self.NN_input_size,))
>         pre_dueling = Dense(512, activation='relu', name='Predueling')(state_input)
>         flatten = Flatten()(pre_dueling)
>         # network separate state value and advantages
>         advantage_fc = Dense(256, activation='relu', name = 'AdvantageFC')(flatten)
>         advantage = Dense(self.action_dim, name = 'Adavantage')(advantage_fc)
>         value_fc = Dense(256, activation='relu', name = 'ValueFC')(flatten)
>         value =  Dense(1, name = 'Value')(value_fc)
>         # network merged and make Q Value
>         Q_values = (value + (advantage - tf.math.reduce_mean(advantage, axis = 1, keepdims = True)))
>         model = Model(inputs=state_input, outputs=Q_values)
>         model.compile(loss=tf.keras.losses.Huber(), optimizer=Adam(self.lr))
>         return model

In a Reinforcement learning settings I simply collect samples and pass them to the below function.

> losses = self.R.model.train_on_batch(self.b_obs, target_q, sample_weight = self.is_weight.flatten())

This worked before on different problem. I am not able to understand when this error can occur