title: Optimization cancelled in Tensorflow
body: <p>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.</p>
  <pre><code>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
  </code></pre>
  <p>`</p>
  <p>Than I do a simple train call as follows</p>
  <pre><code>losses = self.R.model.train_on_batch(self.b_obs, target_q, sample_weight = self.is_weight.flatten())
  </code></pre>
  <p>I get the following error</p>
  <pre><code>W tensorflow/core/data/root_dataset.cc:167] Optimization loop failed:Cancelled:Operation was cancelled.
  </code></pre>