Optimization cancelled in Tensorflow

I am running a NN on keras with tensorflow 2.6.0. The NN is used for RL with a dueling architecture. 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. I have the main DQN model build function below. I simple pass batches to below model and get predictions.

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

`

Than I do a simple train call as follows

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

I get the following error

W tensorflow/core/data/root_dataset.cc:167] Optimization loop failed:Cancelled:Operation was cancelled.