title: keras model.fit() and TF tape.gradient() giving different results
body: <p>I have a model that I am building using the keras functional API. After defining it, I compile it with the SGD optimizer as follows.</p>
  <pre><code>opt = tf.keras.optimizers.SGD(learning_rate=0.05, momentum=0.9, decay=1e-3,clipnorm=1)
  model.compile(optimizer=opt, loss='mse')
  model.fit(train_datagen, epochs=50,shuffle=True,verbose=True)
  </code></pre>
  <p>This works fine and my model converges as expected.</p>
  <p>However, when I to implement the same exact functionality using TF's tape gradient, I consistently get NaN gradients which cause my weights to equate to NaN and subsequently my loss function value because NaN. Here is the code that I use:</p>
  <pre><code>opt = tf.keras.optimizers.SGD(learning_rate=0.05, momentum=0.9, decay=1e-3,clipnorm=1)
  loss_fn = tf.keras.losses.MeanSquaredError()
  
  epochs = 50
  
  for epoch in range(epochs):
          
      batch_list = list(range(len(train_datagen)))
      random.shuffle(batch_list)
  
      running_loss = 0
      
      for ii in batch_list:
          x,y_true = train_datagen[ii]
          
          with tf.GradientTape() as tape:
              y_pred = model(x, training=True)
              loss_value = loss_fn(y_true,y_pred)
              
          grads = tape.gradient(loss_value, model.trainable_variables)
          opt.apply_gradients(zip(grads, model.trainable_variables))
          
          running_loss += loss_value
      
      print('Epoch',epoch,'Running Loss:',running_loss.numpy()/len(batch_list))
  </code></pre>
  <p>Is the code that I wrote equivalent to the Keras model.fit() functionality? For some reason, when I use the above code, I consistently get <code>NaN</code> gradients but with model.fit() it never happens.</p>